Use lib/ts_utils for scripts/sync_protobuf (#12913)

- Simplified `download_file` error handling (400+ return codes already raised errors!)
- Moved `update_metadata` from `scripts/sync_protobuf/_utils.py` to `lib/ts_utils/metadata.py`
- Improved `update_metadata` to support any key (values unvalidated atm) and return the modified dictionary
- Updated `scripts/stubsabot.py` to use `update_metadata`
- Updated `scripts/sync_protobuf/*` to use  `lib/ts_utils`
- Updated `scripts/sync_protobuf/tensorflow.py` and `scripts/sync_protobuf/google_protobuf.py` to use the version directly from the `METADATA.toml` file

scripts/stubsabot.py

@@ -29,7 +29,7 @@ import tomlkit
 from packaging.specifiers import Specifier
 from termcolor import colored
 
-from ts_utils.metadata import StubMetadata, metadata_path, read_metadata
+from ts_utils.metadata import StubMetadata, read_metadata, update_metadata
 from ts_utils.paths import STUBS_PATH, distribution_path
 
 TYPESHED_OWNER = "python"
@@ -653,7 +653,7 @@ _repo_lock = asyncio.Lock()
 BRANCH_PREFIX = "stubsabot"
 
 
-def get_update_pr_body(update: Update, metadata: dict[str, Any]) -> str:
+def get_update_pr_body(update: Update, metadata: Mapping[str, Any]) -> str:
     body = "\n".join(f"{k}: {v}" for k, v in update.links.items())
 
     if update.diff_analysis is not None:
@@ -689,12 +689,7 @@ async def suggest_typeshed_update(update: Update, session: aiohttp.ClientSession
     async with _repo_lock:
         branch_name = f"{BRANCH_PREFIX}/{normalize(update.distribution)}"
         subprocess.check_call(["git", "checkout", "-B", branch_name, "origin/main"])
-        with metadata_path(update.distribution).open("rb") as f:
-            meta = tomlkit.load(f)
-        meta["version"] = update.new_version
-        with metadata_path(update.distribution).open("w", encoding="UTF-8") as f:
-            # tomlkit.dump has partially unknown IO type
-            tomlkit.dump(meta, f)  # pyright: ignore[reportUnknownMemberType]
+        meta = update_metadata(update.distribution, version=update.new_version)
         body = get_update_pr_body(update, meta)
         subprocess.check_call(["git", "commit", "--all", "-m", f"{title}\n\n{body}"])
         if action_level <= ActionLevel.local:
@@ -716,14 +711,9 @@ async def suggest_typeshed_obsolete(obsolete: Obsolete, session: aiohttp.ClientS
     async with _repo_lock:
         branch_name = f"{BRANCH_PREFIX}/{normalize(obsolete.distribution)}"
         subprocess.check_call(["git", "checkout", "-B", branch_name, "origin/main"])
-        with metadata_path(obsolete.distribution).open("rb") as f:
-            meta = tomlkit.load(f)
         obs_string = tomlkit.string(obsolete.obsolete_since_version)
         obs_string.comment(f"Released on {obsolete.obsolete_since_date.date().isoformat()}")
-        meta["obsolete_since"] = obs_string
-        with metadata_path(obsolete.distribution).open("w", encoding="UTF-8") as f:
-            # tomlkit.dump has partially unknown Mapping type
-            tomlkit.dump(meta, f)  # pyright: ignore[reportUnknownMemberType]
+        update_metadata(obsolete.distribution, obsolete_since=obs_string)
         body = "\n".join(f"{k}: {v}" for k, v in obsolete.links.items())
         subprocess.check_call(["git", "commit", "--all", "-m", f"{title}\n\n{body}"])
         if action_level <= ActionLevel.local:

scripts/sync_protobuf/_utils.py

@@ -3,12 +3,10 @@ from __future__ import annotations
 import subprocess
 import sys
 from http.client import HTTPResponse
-from pathlib import Path
 from typing import TYPE_CHECKING, Iterable
 from urllib.request import urlopen
 from zipfile import ZipFile
 
-import tomlkit
 from mypy_protobuf.main import (  # type: ignore[import-untyped]  # pyright: ignore[reportMissingTypeStubs]
     __version__ as mypy_protobuf__version__,
 )
@@ -16,18 +14,14 @@ from mypy_protobuf.main import (  # type: ignore[import-untyped]  # pyright: ign
 if TYPE_CHECKING:
     from _typeshed import StrOrBytesPath, StrPath
 
-REPO_ROOT = Path(__file__).absolute().parent.parent.parent
 MYPY_PROTOBUF_VERSION = mypy_protobuf__version__
 
 
 def download_file(url: str, destination: StrPath) -> None:
     print(f"Downloading '{url}' to '{destination}'")
     resp: HTTPResponse
-    with urlopen(url) as resp:
-        if resp.getcode() != 200:
-            raise RuntimeError(f"Error downloading {url}")
-        with open(destination, "wb") as file:
-            file.write(resp.read())
+    with urlopen(url) as resp, open(destination, "wb") as file:
+        file.write(resp.read())
 
 
 def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
@@ -36,17 +30,6 @@ def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
         file_in.extractall(destination)
 
 
-def update_metadata(metadata_folder: StrPath, new_extra_description: str) -> None:
-    metadata_path = Path(metadata_folder) / "METADATA.toml"
-    with open(metadata_path) as file:
-        metadata = tomlkit.load(file)
-    metadata["extra_description"] = new_extra_description
-    with open(metadata_path, "w") as file:
-        # tomlkit.dump has partially unknown IO type
-        tomlkit.dump(metadata, file)  # pyright: ignore[reportUnknownMemberType]
-    print(f"Updated {metadata_path}")
-
-
 def run_protoc(
     proto_paths: Iterable[StrPath], mypy_out: StrPath, proto_globs: Iterable[str], cwd: StrOrBytesPath | None = None
 ) -> str:

scripts/sync_protobuf/google_protobuf.py

@@ -14,13 +14,14 @@ import tempfile
 from pathlib import Path
 from typing import Any
 
-from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+from _utils import MYPY_PROTOBUF_VERSION, download_file, extract_archive, run_protoc
+from ts_utils.metadata import read_metadata, update_metadata
+from ts_utils.paths import distribution_path
 
-# Whenever you update PACKAGE_VERSION here, version should be updated
-# in stubs/protobuf/METADATA.toml and vice-versa.
-PACKAGE_VERSION = "28.2"
+# PyPi version has an extra "major" number that the Git version doesn't have
+PACKAGE_VERSION = read_metadata("protobuf").version_spec.version[2:]
 
-STUBS_FOLDER = REPO_ROOT / "stubs" / "protobuf"
+STUBS_FOLDER = distribution_path("protobuf").absolute()
 ARCHIVE_FILENAME = f"protobuf-{PACKAGE_VERSION}.zip"
 ARCHIVE_URL = f"https://github.com/protocolbuffers/protobuf/releases/download/v{PACKAGE_VERSION}/{ARCHIVE_FILENAME}"
 EXTRACTED_PACKAGE_DIR = f"protobuf-{PACKAGE_VERSION}"
@@ -78,13 +79,14 @@ def main() -> None:
     shutil.rmtree(temp_dir)
 
     update_metadata(
-        STUBS_FOLDER,
-        f"""Partially generated using \
+        "protobuf",
+        extra_description=f"""Partially generated using \
 [mypy-protobuf=={MYPY_PROTOBUF_VERSION}](https://github.com/nipunn1313/mypy-protobuf/tree/v{MYPY_PROTOBUF_VERSION}) \
 and {PROTOC_VERSION} on \
 [protobuf v{PACKAGE_VERSION}](https://github.com/protocolbuffers/protobuf/releases/tag/v{PACKAGE_VERSION}) \
 (python `protobuf=={PYTHON_PROTOBUF_VERSION}`).""",
     )
+    print("Updated protobuf/METADATA.toml")
 
     # Run pre-commit to cleanup the stubs
     subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))

scripts/sync_protobuf/s2clientprotocol.py

@@ -12,13 +12,15 @@ import sys
 import tempfile
 from pathlib import Path
 
-from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+from _utils import MYPY_PROTOBUF_VERSION, download_file, extract_archive, run_protoc
+from ts_utils.metadata import update_metadata
+from ts_utils.paths import distribution_path
 
 # Whenever you update PACKAGE_VERSION here, version should be updated
 # in stubs/s2clientprotocol/METADATA.toml and vice-versa.
 PACKAGE_VERSION = "c04df4adbe274858a4eb8417175ee32ad02fd609"
 
-STUBS_FOLDER = REPO_ROOT / "stubs" / "s2clientprotocol"
+STUBS_FOLDER = distribution_path("s2clientprotocol").absolute()
 ARCHIVE_FILENAME = f"{PACKAGE_VERSION}.zip"
 ARCHIVE_URL = f"https://github.com/Blizzard/s2client-proto/archive/{ARCHIVE_FILENAME}"
 EXTRACTED_PACKAGE_DIR = f"s2client-proto-{PACKAGE_VERSION}"
@@ -57,12 +59,13 @@ def main() -> None:
     shutil.rmtree(temp_dir)
 
     update_metadata(
-        STUBS_FOLDER,
-        f"""Partially generated using \
+        "s2clientprotocol",
+        extra_description=f"""Partially generated using \
 [mypy-protobuf=={MYPY_PROTOBUF_VERSION}](https://github.com/nipunn1313/mypy-protobuf/tree/v{MYPY_PROTOBUF_VERSION}) \
 and {PROTOC_VERSION} on \
 [s2client-proto {PYTHON_S2_CLIENT_PROTO_VERSION}](https://github.com/Blizzard/s2client-proto/tree/{PACKAGE_VERSION}).""",
     )
+    print("Updated s2clientprotocol/METADATA.toml")
 
     # Run pre-commit to cleanup the stubs
     subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))

scripts/sync_protobuf/tensorflow.py

@@ -13,13 +13,13 @@ import sys
 import tempfile
 from pathlib import Path
 
-from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+from _utils import MYPY_PROTOBUF_VERSION, download_file, extract_archive, run_protoc
+from ts_utils.metadata import read_metadata, update_metadata
+from ts_utils.paths import distribution_path
 
-# Whenever you update PACKAGE_VERSION here, version should be updated
-# in stubs/tensorflow/METADATA.toml and vice-versa.
-PACKAGE_VERSION = "2.17.0"
+PACKAGE_VERSION = read_metadata("tensorflow").version_spec.version
 
-STUBS_FOLDER = REPO_ROOT / "stubs" / "tensorflow"
+STUBS_FOLDER = distribution_path("tensorflow").absolute()
 ARCHIVE_FILENAME = f"v{PACKAGE_VERSION}.zip"
 ARCHIVE_URL = f"https://github.com/tensorflow/tensorflow/archive/refs/tags/{ARCHIVE_FILENAME}"
 EXTRACTED_PACKAGE_DIR = f"tensorflow-{PACKAGE_VERSION}"
@@ -123,11 +123,12 @@ def main() -> None:
     post_creation()
 
     update_metadata(
-        STUBS_FOLDER,
-        f"""Partially generated using \
+        "tensorflow",
+        extra_description=f"""Partially generated using \
 [mypy-protobuf=={MYPY_PROTOBUF_VERSION}](https://github.com/nipunn1313/mypy-protobuf/tree/v{MYPY_PROTOBUF_VERSION}) \
 and {PROTOC_VERSION} on `tensorflow=={PACKAGE_VERSION}`.""",
     )
+    print("Updated tensorflow/METADATA.toml")
 
     # Run pre-commit to cleanup the stubs
     subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))