Rewrote protobuf generation scripts in Python (#12527)

requirements-tests.txt

@@ -14,6 +14,10 @@ ruff==0.5.4              # must match .pre-commit-config.yaml
 
 # Libraries used by our various scripts.
 aiohttp==3.10.2
+# grpc install only fails on Windows, but let's avoid building sdist on other platforms
+# https://github.com/grpc/grpc/issues/36201
+grpcio-tools; python_version < "3.13"
+mypy-protobuf==3.6.0
 packaging==24.1
 pathspec>=0.11.1
 pre-commit

scripts/sync_protobuf/_utils.py

@@ -0,0 +1,67 @@
+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__,
+)
+
+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())
+
+
+def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
+    print(f"Extracting '{archive_path}' to '{destination}'")
+    with ZipFile(archive_path) as file_in:
+        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:
+    """TODO: Describe parameters and return"""
+    protoc_version = (
+        subprocess.run([sys.executable, "-m", "grpc_tools.protoc", "--version"], capture_output=True).stdout.decode().strip()
+    )
+    print()
+    print(protoc_version)
+    protoc_args = [
+        *[f"--proto_path={proto_path}" for proto_path in proto_paths],
+        "--mypy_out",
+        f"relax_strict_optional_primitives:{mypy_out}",
+        *proto_globs,
+    ]
+    print("Running: protoc\n    " + "\n    ".join(protoc_args) + "\n")
+    subprocess.run((sys.executable, "-m", "grpc_tools.protoc", *protoc_args), cwd=cwd, check=True)
+    return protoc_version

scripts/sync_protobuf/google_protobuf.py

@@ -0,0 +1,94 @@
+"""
+Generates the protobuf stubs for the given protobuf version using mypy-protobuf.
+Generally, new minor versions are a good time to update the stubs.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+import sys
+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
+
+# Whenever you update PACKAGE_VERSION here, version should be updated
+# in stubs/protobuf/METADATA.toml and vice-versa.
+PACKAGE_VERSION = "27.1"
+
+STUBS_FOLDER = REPO_ROOT / "stubs" / "protobuf"
+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}"
+
+VERSION_PATTERN = re.compile(r'def game_version\(\):\n    return "(.+?)"')
+PROTO_FILE_PATTERN = re.compile(r'"//:(.*)_proto"')
+
+
+def extract_python_version(file_path: Path) -> str:
+    """Extract the Python version from https://github.com/protocolbuffers/protobuf/blob/main/version.json"""
+    with open(file_path) as file:
+        data: dict[str, Any] = json.load(file)
+    # The root key will be the protobuf source code version
+    version = next(iter(data.values()))["languages"]["python"]
+    assert isinstance(version, str)
+    return version
+
+
+def extract_proto_file_paths(temp_dir: Path) -> list[str]:
+    """
+    Roughly reproduce the subset of .proto files on the public interface
+    as described in py_proto_library calls in
+    https://github.com/protocolbuffers/protobuf/blob/main/python/dist/BUILD.bazel
+    """
+    with open(temp_dir / EXTRACTED_PACKAGE_DIR / "python" / "dist" / "BUILD.bazel") as file:
+        matched_lines = filter(None, (re.search(PROTO_FILE_PATTERN, line) for line in file))
+        proto_files = [
+            EXTRACTED_PACKAGE_DIR + "/src/google/protobuf/" + match.group(1).replace("compiler_", "compiler/") + ".proto"
+            for match in matched_lines
+        ]
+    return proto_files
+
+
+def main() -> None:
+    temp_dir = Path(tempfile.mkdtemp())
+    # Fetch s2clientprotocol (which contains all the .proto files)
+    archive_path = temp_dir / ARCHIVE_FILENAME
+    download_file(ARCHIVE_URL, archive_path)
+    extract_archive(archive_path, temp_dir)
+
+    # Remove existing pyi
+    for old_stub in STUBS_FOLDER.rglob("*_pb2.pyi"):
+        old_stub.unlink()
+
+    PROTOC_VERSION = run_protoc(
+        proto_paths=(f"{EXTRACTED_PACKAGE_DIR}/src",),
+        mypy_out=STUBS_FOLDER,
+        proto_globs=extract_proto_file_paths(temp_dir),
+        cwd=temp_dir,
+    )
+
+    PYTHON_PROTOBUF_VERSION = extract_python_version(temp_dir / EXTRACTED_PACKAGE_DIR / "version.json")
+
+    # Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
+    shutil.rmtree(temp_dir)
+
+    update_metadata(
+        STUBS_FOLDER,
+        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}`).""",
+    )
+
+    # Run pre-commit to cleanup the stubs
+    subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
+
+
+if __name__ == "__main__":
+    main()

scripts/sync_protobuf/s2clientprotocol.py

@@ -0,0 +1,72 @@
+"""
+Generates the protobuf stubs for the given s2clientprotocol version using mypy-protobuf.
+Generally, new minor versions are a good time to update the stubs.
+"""
+
+from __future__ import annotations
+
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+
+# 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"
+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}"
+
+VERSION_PATTERN = re.compile(r'def game_version\(\):\n    return "(.+?)"')
+
+
+def extract_python_version(file_path: Path) -> str:
+    """Extract Python version from s2clientprotocol's build file"""
+    match = re.search(VERSION_PATTERN, file_path.read_text())
+    assert match
+    return match.group(1)
+
+
+def main() -> None:
+    temp_dir = Path(tempfile.mkdtemp())
+    # Fetch s2clientprotocol (which contains all the .proto files)
+    archive_path = temp_dir / ARCHIVE_FILENAME
+    download_file(ARCHIVE_URL, archive_path)
+    extract_archive(archive_path, temp_dir)
+
+    # Remove existing pyi
+    for old_stub in STUBS_FOLDER.rglob("*_pb2.pyi"):
+        old_stub.unlink()
+
+    PROTOC_VERSION = run_protoc(
+        proto_paths=(EXTRACTED_PACKAGE_DIR,),
+        mypy_out=STUBS_FOLDER,
+        proto_globs=(f"{EXTRACTED_PACKAGE_DIR}/s2clientprotocol/*.proto",),
+        cwd=temp_dir,
+    )
+
+    PYTHON_S2_CLIENT_PROTO_VERSION = extract_python_version(temp_dir / EXTRACTED_PACKAGE_DIR / "s2clientprotocol" / "build.py")
+
+    # Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
+    shutil.rmtree(temp_dir)
+
+    update_metadata(
+        STUBS_FOLDER,
+        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}).""",
+    )
+
+    # Run pre-commit to cleanup the stubs
+    subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
+
+
+if __name__ == "__main__":
+    main()