stubsabot updates

hauntsaninja · hauntsaninja · commit 2cc1c9b5c6a1 · 2022-06-09T01:54:54.000-07:00
diff --git a/scripts/stubsabot.py b/scripts/stubsabot.py
@@ -4,13 +4,18 @@
 import argparse
 import asyncio
 import datetime
+import enum
+import io
 import os
 import re
 import subprocess
 import sys
+import tarfile
 import urllib.parse
+import zipfile
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Any
 
 import aiohttp
 import packaging.specifiers
@@ -19,6 +24,12 @@
 import tomlkit
 
 
+class ActionLevel(enum.IntEnum):
+    nothing = 0  # make no changes
+    local = 1  # make changes that affect local repo
+    everything = 2  # do everything, e.g. open PRs
+
+
 @dataclass
 class StubInfo:
     distribution: str
@@ -43,6 +54,7 @@ class PypiInfo:
     distribution: str
     version: packaging.version.Version
     upload_date: datetime.datetime
+    release_to_download: dict[str, Any]
 
 
 async def fetch_pypi_info(distribution: str, session: aiohttp.ClientSession) -> PypiInfo:
@@ -51,9 +63,14 @@ async def fetch_pypi_info(distribution: str, session: aiohttp.ClientSession) ->
         response.raise_for_status()
         j = await response.json()
         version = j["info"]["version"]
-        date = datetime.datetime.fromisoformat(j["releases"][version][0]["upload_time"])
+        # prefer wheels, since it's what most users will get / it's pretty easy to mess up MANIFEST
+        release_to_download = sorted(j["releases"][version], key=lambda x: bool(x["packagetype"] == "bdist_wheel"))[-1]
+        date = datetime.datetime.fromisoformat(release_to_download["upload_time"])
         return PypiInfo(
-            distribution=distribution, version=packaging.version.Version(version), upload_date=date
+            distribution=distribution,
+            version=packaging.version.Version(version),
+            upload_date=date,
+            release_to_download=release_to_download,
         )
 
 
@@ -64,17 +81,47 @@ class Update:
     old_version_spec: str
     new_version_spec: str
 
+    def __str__(self) -> str:
+        return f"Updating {self.distribution} from {self.old_version_spec!r} to {self.new_version_spec!r}"
+
+
+@dataclass
+class Obsolete:
+    distribution: str
+    stub_path: Path
+    obsolete_since_version: str
+
+    def __str__(self) -> str:
+        return f"Marking {self.distribution} as obsolete since {self.obsolete_since_version!r}"
+
 
 @dataclass
 class NoUpdate:
     distribution: str
     reason: str
 
+    def __str__(self) -> str:
+        return f"Skipping {self.distribution}: {self.reason}"
+
+
+async def package_contains_py_typed(release_to_download: dict[str, Any], session: aiohttp.ClientSession) -> bool:
+    async with session.get(release_to_download["url"]) as response:
+        body = io.BytesIO(await response.read())
+
+    if release_to_download["packagetype"] == "bdist_wheel":
+        assert release_to_download["filename"].endswith(".whl")
+        with zipfile.ZipFile(body) as zf:
+            return any(Path(f).name == "py.typed" for f in zf.namelist())
+    elif release_to_download["packagetype"] == "sdist":
+        assert release_to_download["filename"].endswith(".tar.gz")
+        with tarfile.open(fileobj=body, mode="r:gz") as zf:
+            return any(Path(f).name == "py.typed" for f in zf.getnames())
+    else:
+        raise AssertionError
+
 
 def _check_spec(updated_spec: str, version: packaging.version.Version) -> str:
-    assert version in packaging.specifiers.SpecifierSet(
-        "==" + updated_spec
-    ), f"{version} not in {updated_spec}"
+    assert version in packaging.specifiers.SpecifierSet("==" + updated_spec), f"{version} not in {updated_spec}"
     return updated_spec
 
 
@@ -89,7 +136,7 @@ def get_updated_version_spec(spec: str, version: packaging.version.Version) -> s
     return _check_spec(".".join(rounded_version) + ".*", version)
 
 
-async def determine_action(stub_path: Path, session: aiohttp.ClientSession) -> Update | NoUpdate:
+async def determine_action(stub_path: Path, session: aiohttp.ClientSession) -> Update | NoUpdate | Obsolete:
     stub_info = read_typeshed_stub_metadata(stub_path)
     if stub_info.obsolete:
         return NoUpdate(stub_info.distribution, "obsolete")
@@ -101,6 +148,9 @@ async def determine_action(stub_path: Path, session: aiohttp.ClientSession) -> U
     if pypi_info.version in spec:
         return NoUpdate(stub_info.distribution, "up to date")
 
+    if await package_contains_py_typed(pypi_info.release_to_download, session):
+        return Obsolete(stub_info.distribution, stub_path, obsolete_since_version=str(pypi_info.version))
+
     return Update(
         distribution=stub_info.distribution,
         stub_path=stub_path,
@@ -109,23 +159,60 @@ async def determine_action(stub_path: Path, session: aiohttp.ClientSession) -> U
     )
 
 
+TYPESHED_OWNER = "python"
+FORK_OWNER = "hauntsaninja"
+
+
+async def create_or_update_pull_request(title: str, branch_name: str, session: aiohttp.ClientSession):
+    secret = os.environ["GITHUB_TOKEN"]
+    if secret.startswith("ghp"):
+        auth = f"token {secret}"
+    else:
+        auth = f"Bearer {secret}"
+
+    async with session.post(
+        f"https://api.github.com/repos/{TYPESHED_OWNER}/typeshed/pulls",
+        json={"title": title, "head": f"{FORK_OWNER}:{branch_name}", "base": "master"},
+        headers={"Accept": "application/vnd.github.v3+json", "Authorization": auth},
+    ) as response:
+        body = await response.json()
+        if response.status == 422 and any(
+            "A pull request already exists" in e.get("message", "") for e in body.get("errors", [])
+        ):
+            # Find the existing PR
+            async with session.get(
+                f"https://api.github.com/repos/{TYPESHED_OWNER}/typeshed/pulls",
+                params={"state": "open", "head": f"{FORK_OWNER}:{branch_name}", "base": "master"},
+                headers={"Accept": "application/vnd.github.v3+json", "Authorization": auth},
+            ) as response:
+                response.raise_for_status()
+                body = await response.json()
+                assert len(body) >= 1
+                pr_number = body[0]["number"]
+            # Update the PR's title
+            async with session.patch(
+                f"https://api.github.com/repos/{TYPESHED_OWNER}/typeshed/pulls/{pr_number}",
+                json={"title": title},
+                headers={"Accept": "application/vnd.github.v3+json", "Authorization": auth},
+            ) as response:
+                response.raise_for_status()
+            return
+        response.raise_for_status()
+
+
 def normalize(name: str) -> str:
     # PEP 503 normalization
     return re.sub(r"[-_.]+", "-", name).lower()
 
 
+# lock should be unnecessary, but can't hurt to enforce mutual exclusion
 _repo_lock = asyncio.Lock()
 
-TYPESHED_OWNER = "python"
-FORK_OWNER = "hauntsaninja"
-
 
-async def suggest_typeshed_update(
-    update: Update, session: aiohttp.ClientSession, dry_run: bool
-) -> None:
+async def suggest_typeshed_update(update: Update, session: aiohttp.ClientSession, action_level: ActionLevel) -> None:
+    if action_level <= ActionLevel.nothing:
+        return
     title = f"[stubsabot] Bump {update.distribution} to {update.new_version_spec}"
-
-    # lock should be unnecessary, but can't hurt to enforce mutual exclusion
     async with _repo_lock:
         branch_name = f"stubsabot/{normalize(update.distribution)}"
         subprocess.check_call(["git", "checkout", "-B", branch_name, "origin/master"])
@@ -135,56 +222,66 @@ async def suggest_typeshed_update(
         with open(update.stub_path / "METADATA.toml", "w") as f:
             tomlkit.dump(meta, f)
         subprocess.check_call(["git", "commit", "--all", "-m", title])
-        if dry_run:
+        if action_level <= ActionLevel.local:
             return
         subprocess.check_call(["git", "push", "origin", branch_name, "--force-with-lease"])
 
-    secret = os.environ["GITHUB_TOKEN"]
-    if secret.startswith("ghp"):
-        auth = f"token {secret}"
-    else:
-        auth = f"Bearer {secret}"
+    await create_or_update_pull_request(title, branch_name, session)
 
-    async with session.post(
-        f"https://api.github.com/repos/{TYPESHED_OWNER}/typeshed/pulls",
-        json={"title": title, "head": f"{FORK_OWNER}:{branch_name}", "base": "master"},
-        headers={"Accept": "application/vnd.github.v3+json", "Authorization": auth},
-    ) as response:
-        body = await response.json()
-        if response.status == 422 and any(
-            "A pull request already exists" in e.get("message", "") for e in body.get("errors", [])
-        ):
-            # TODO: diff and update existing pull request
+
+async def suggest_typeshed_obsolete(obsolete: Obsolete, session: aiohttp.ClientSession, action_level: ActionLevel) -> None:
+    if action_level <= ActionLevel.nothing:
+        return
+    title = f"[stubsabot] Mark {obsolete.distribution} as obsolete since {obsolete.obsolete_since_version}"
+    async with _repo_lock:
+        branch_name = f"stubsabot/{normalize(obsolete.distribution)}"
+        subprocess.check_call(["git", "checkout", "-B", branch_name, "origin/master"])
+        with open(obsolete.stub_path / "METADATA.toml", "rb") as f:
+            meta = tomlkit.load(f)
+        meta["obsolete_since"] = obsolete.obsolete_since_version
+        with open(obsolete.stub_path / "METADATA.toml", "w") as f:
+            tomlkit.dump(meta, f)
+        subprocess.check_call(["git", "commit", "--all", "-m", title])
+        if action_level <= ActionLevel.local:
             return
-        response.raise_for_status()
+        subprocess.check_call(["git", "push", "origin", branch_name, "--force-with-lease"])
+
+    await create_or_update_pull_request(title, branch_name, session)
 
 
 async def main() -> None:
     assert sys.version_info >= (3, 9)
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument(
+        "--action-level",
+        type=lambda x: getattr(ActionLevel, x),  # type: ignore[no-any-return]
+        default=ActionLevel.everything,
+        help="Limit actions performed to achieve dry runs for different levels of dryness",
+    )
     args = parser.parse_args()
 
     try:
         conn = aiohttp.TCPConnector(limit_per_host=10)
         async with aiohttp.ClientSession(connector=conn) as session:
-            tasks = [
-                asyncio.create_task(determine_action(stubs_path, session))
-                for stubs_path in Path("stubs").iterdir()
-            ]
+            tasks = [asyncio.create_task(determine_action(stubs_path, session)) for stubs_path in Path("stubs").iterdir()]
             for task in asyncio.as_completed(tasks):
                 update = await task
+                print(update)
                 if isinstance(update, NoUpdate):
                     continue
                 if isinstance(update, Update):
-                    await suggest_typeshed_update(update, session, dry_run=args.dry_run)
+                    await suggest_typeshed_update(update, session, action_level=args.action_level)
+                    continue
+                if isinstance(update, Obsolete):
+                    await suggest_typeshed_obsolete(update, session, action_level=args.action_level)
                     continue
                 raise AssertionError
     finally:
         # if you need to cleanup, try:
         # git branch -D $(git branch --list 'stubsabot/*')
-        subprocess.check_call(["git", "checkout", "master"])
+        if args.action_level >= ActionLevel.local:
+            subprocess.check_call(["git", "checkout", "master"])
 
 
 if __name__ == "__main__":
