scripts/collect_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import shutil\n",
    "import subprocess\n",
    "import time\n",
    "from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed\n",
    "from pathlib import Path\n",
    "\n",
    "import libcst as cst\n",
    "from tqdm import tqdm\n",
    "\n",
    "from typet5.data import GitRepo, get_dataset_dir\n",
    "from typet5.type_env import collect_annots_info, mypy_checker\n",
    "from typet5.utils import proj_root, read_file, write_file, not_none\n",
    "\n",
    "os.chdir(proj_root())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Repos already downloaded.\n",
      "Reading last updates...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4890/4890 [00:27<00:00, 175.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloaded 4890/5996 repos.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# download all candidate repos\n",
    "\n",
    "all_repos = json.loads(read_file(\"data/mypy-dependents-by-stars.json\"))\n",
    "all_repos = [GitRepo.from_json(r) for r in all_repos]\n",
    "# all_repos=all_repos[:10] # for testing\n",
    "\n",
    "repos_dir = get_dataset_dir(\"ManyTypes4Py\") / \"repos\"\n",
    "\n",
    "def clear_downloaded_repos(repos_dir):\n",
    "    shutil.rmtree(repos_dir)\n",
    "\n",
    "\n",
    "def download_repos(\n",
    "    to_download: list[GitRepo], repos_dir, download_timeout=10.0, max_workers=10\n",
    ") -> list[GitRepo]:\n",
    "    def download_single(repo: GitRepo):\n",
    "        try:\n",
    "            if repo.download(repos_dir, timeout=download_timeout):\n",
    "                repo.read_last_update(repos_dir)\n",
    "                return repo\n",
    "            else:\n",
    "                return None\n",
    "        except subprocess.TimeoutExpired:\n",
    "            return None\n",
    "        except Exception as e:\n",
    "            logging.warning(f\"Failed to download {repo.name}. Exception: {e}\")\n",
    "            return None\n",
    "\n",
    "    print(\"Downloading repos from Github...\")\n",
    "    t_start = time.time()\n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "        fs = [executor.submit(download_single, repo) for repo in to_download]\n",
    "        rs = [f.result() for f in tqdm(as_completed(fs), total=len(fs))]\n",
    "    print(f\"Downloading took {time.time() - t_start} seconds.\")\n",
    "    downloaded = [r for r in rs if r is not None]\n",
    "    return downloaded\n",
    "\n",
    "\n",
    "if not repos_dir.exists():\n",
    "    (repos_dir / \"downloading\").mkdir(parents=True)\n",
    "    (repos_dir / \"downloaded\").mkdir(parents=True)\n",
    "    downloaded_repos = download_repos(all_repos, repos_dir)\n",
    "    print(\"Deleting failed repos...\")\n",
    "    shutil.rmtree(repos_dir / \"downloading\")\n",
    "else:\n",
    "    print(\"Repos already downloaded.\")\n",
    "    downloaded_dirs = set(d.name for d in (repos_dir / \"downloaded\").iterdir())\n",
    "    downloaded_repos = [r for r in all_repos if r.authorname() in downloaded_dirs]\n",
    "    print(\"Reading last updates...\")\n",
    "    for r in tqdm(downloaded_repos):\n",
    "        r.read_last_update(repos_dir)\n",
    "print(f\"Downloaded {len(downloaded_repos)}/{len(all_repos)} repos.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1218 / 4890 repos are updated within a year.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1218/1218 [00:05<00:00, 243.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1181/1218 repos are within the size limit.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# filter out repos that are too old or too big\n",
    "\n",
    "from datetime import datetime, timezone\n",
    "\n",
    "date_threshold = datetime(2021, 4, 20)\n",
    "new_repos = [r for r in downloaded_repos if not_none(r.last_update) > date_threshold]\n",
    "print(f\"{len(new_repos)} / {len(downloaded_repos)} repos are updated within a year.\")\n",
    "loc_limit = 50000\n",
    "\n",
    "small_repos = []\n",
    "for rep in tqdm(new_repos):\n",
    "    try:\n",
    "        loc = rep.count_lines_of_code(repos_dir)\n",
    "        if loc < loc_limit:\n",
    "            small_repos.append(rep)\n",
    "    except UnicodeDecodeError:\n",
    "        # nothing we can do\n",
    "        pass\n",
    "    except Exception as e:\n",
    "        logging.warning(f\"Failed to count lines of code for {rep.name}. Exception: {e}\")\n",
    "\n",
    "print(\n",
    "    f\"{len(small_repos)}/{len(new_repos)} repos are within the size limit ({loc_limit} LOC).\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter away repos with too few annotations\n",
    "\n",
    "def count_repo_annots(rep):\n",
    "    try:\n",
    "        rep.count_annotations(repos_dir)\n",
    "        if rep.n_type_annots / rep.lines_of_code > 0.05:\n",
    "            return rep\n",
    "    except Exception as e:\n",
    "        logging.warning(f\"Failed to count annotations for {rep.name}. Exception: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "with ProcessPoolExecutor(max_workers=30) as executor:\n",
    "    fs = [executor.submit(count_repo_annots, rep) for rep in small_repos]\n",
    "    rs = [f.result() for f in tqdm(as_completed(fs), total=len(fs))]\n",
    "useful_repos: list[GitRepo] = [\n",
    "    r for r in rs if r is not None and \"typeshed\" not in r.name\n",
    "]\n",
    "\n",
    "print(\n",
    "    f\"{len(useful_repos)}/{len(small_repos)} repos are parsable and have enough portions of type annotations.\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of manual annotations: 343595\n",
      "Total number of type places: 544497\n",
      "Total number of lines of code: 3342911\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[GitRepo(author='skorokithakis', name='catt', url='https://github.com/skorokithakis/catt', stars=1740, forks=762, lines_of_code=2036, last_update=datetime.datetime(2022, 4, 10, 1, 30, 43), n_type_annots=140, n_type_places=433),\n",
       " GitRepo(author='encode', name='databases', url='https://github.com/encode/databases', stars=769, forks=48, lines_of_code=3124, last_update=datetime.datetime(2022, 3, 6, 12, 25, 10), n_type_annots=323, n_type_places=498),\n",
       " GitRepo(author='Curt-Park', name='rainbow-is-all-you-need', url='https://github.com/Curt-Park/rainbow-is-all-you-need', stars=490, forks=110, lines_of_code=107, last_update=datetime.datetime(2022, 1, 13, 23, 4, 48), n_type_annots=26, n_type_places=30),\n",
       " GitRepo(author='jreese', name='aiomultiprocess', url='https://github.com/jreese/aiomultiprocess', stars=585, forks=45, lines_of_code=1140, last_update=datetime.datetime(2022, 2, 4, 21, 28, 7), n_type_annots=138, n_type_places=213),\n",
       " GitRepo(author='instaloader', name='instaloader', url='https://github.com/instaloader/instaloader', stars=874, forks=134, lines_of_code=5417, last_update=datetime.datetime(2022, 4, 18, 9, 49, 34), n_type_annots=569, n_type_places=843)]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Some summary statistics\n",
    "\n",
    "# print total number of manual annotations\n",
    "n_total_annots = sum(not_none(rep.n_type_annots) for rep in useful_repos)\n",
    "print(\"Total number of manual annotations:\", n_total_annots)\n",
    "\n",
    "# print total number of type places\n",
    "n_total_places = sum(not_none(rep.n_type_places) for rep in useful_repos)\n",
    "print(\"Total number of type places:\", n_total_places)\n",
    "\n",
    "# print total number of lines of code\n",
    "n_total_lines = sum(not_none(rep.lines_of_code) for rep in useful_repos)\n",
    "print(\"Total number of lines of code:\", n_total_lines)\n",
    "\n",
    "# print average number of type annotations per line of code excluding projects with more than 1000 lines of code\n",
    "n_avg_annots = (\n",
    "    sum(not_none(rep.n_type_annots) for rep in useful_repos if rep.lines_of_code < 1000)\n",
    "    / n_total_lines\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[GitRepo(author='typeddjango', name='pytest-mypy-plugins', url='https://github.com/typeddjango/pytest-mypy-plugins', stars=12, forks=0, lines_of_code=1039, last_update=datetime.datetime(2022, 4, 18, 23, 25, 40), n_type_annots=155, n_type_places=158), GitRepo(author='jfly', name='jfly.github.io', url='https://github.com/jfly/jfly.github.io', stars=0, forks=0, lines_of_code=650, last_update=datetime.datetime(2022, 4, 12, 8, 23, 39), n_type_annots=39, n_type_places=122), GitRepo(author='seattleflu', name='id3c', url='https://github.com/seattleflu/id3c', stars=2, forks=0, lines_of_code=8883, last_update=datetime.datetime(2022, 4, 21, 15, 38, 59), n_type_annots=675, n_type_places=1068)]\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "\n",
    "useful_repos_path = proj_root() / \"scripts\" / \"useful_repos.pkl\"\n",
    "with useful_repos_path.open(\"wb\") as f:\n",
    "    pickle.dump(useful_repos, f)\n",
    "print(f\"Saved {len(useful_repos)} useful repos to {useful_repos_path}.\")\n",
    "with useful_repos_path.open(\"rb\") as f:\n",
    "    print(pickle.load(f)[:3])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below tries to split the dataset based on the original dataset split used by the paper. But since that the list of repos returned by the GitHub query above can change over time, some repos might no longer be present. You might consider perform your own splitting if the issue is serious."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typet5.utils import pickle_load, Path, proj_root, tqdm\n",
    "import shutil\n",
    "from typet5.data import GitRepo, get_dataset_dir\n",
    "\n",
    "os.chdir(proj_root())\n",
    "\n",
    "repos_split = pickle_load(Path(\"data/repos_split.pkl\"))\n",
    "repos_dir = get_dataset_dir(\"ManyTypes4Py\") / \"repos\"\n",
    "\n",
    "for split, repos in repos_split.items():\n",
    "    for r in tqdm(repos, desc=f\"Moving {split} repos.\"):\n",
    "        r: GitRepo\n",
    "        split: str\n",
    "        src = repos_dir / r.authorname()\n",
    "        (repos_dir / split).mkdir(parents=True, exist_ok=True)\n",
    "        dest = repos_dir / split / r.authorname()\n",
    "        if src.exists():\n",
    "            shutil.move(src, dest)\n",
    "        else:\n",
    "            print(f\"Repo {r.name} not found.\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f6ffc72953da4dd16b2e00785be9c4013ef131f465a8658f3921b6634d4eeec8"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('.venv': pipenv)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}