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/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")))