From db53bdada524e4f737fdd38d8cd7cd350eb8aa13 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Thu, 15 Aug 2024 02:06:32 -0400
Subject: [PATCH 1/9] WIP (s2clientprotocol done)

---
 requirements-tests.txt               |  3 ++
 scripts/protobuf/__init__.py         |  0
 scripts/protobuf/_helpers.py         | 55 +++++++++++++++++++++
 scripts/protobuf/s2clientprotocol.py | 73 ++++++++++++++++++++++++++++
 stubs/s2clientprotocol/METADATA.toml |  2 +-
 5 files changed, 132 insertions(+), 1 deletion(-)
 create mode 100644 scripts/protobuf/__init__.py
 create mode 100644 scripts/protobuf/_helpers.py
 create mode 100644 scripts/protobuf/s2clientprotocol.py

diff --git a/requirements-tests.txt b/requirements-tests.txt
index 89b009d30caf..99e05ecfe5a8 100644
--- a/requirements-tests.txt
+++ b/requirements-tests.txt
@@ -13,9 +13,12 @@ ruff==0.5.4              # must match .pre-commit-config.yaml
 
 # Libraries used by our various scripts.
 aiohttp==3.10.2
+grpcio-tools
+mypy-protobuf==3.6.0
 packaging==24.1
 pathspec>=0.11.1
 pre-commit
+requests
 stubdefaulter==0.1.0
 termcolor>=2.3
 tomli==2.0.1
diff --git a/scripts/protobuf/__init__.py b/scripts/protobuf/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/scripts/protobuf/_helpers.py b/scripts/protobuf/_helpers.py
new file mode 100644
index 000000000000..f7debe9ac9e7
--- /dev/null
+++ b/scripts/protobuf/_helpers.py
@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+import zipfile
+from typing import TYPE_CHECKING, Iterable
+
+import requests
+import tomlkit
+
+if TYPE_CHECKING:
+    from _typeshed import FileDescriptorOrPath, StrOrBytesPath, StrPath
+
+
+def download_file(url: str | bytes, destination: FileDescriptorOrPath) -> None:
+    print(f"Downloading {url!r} to '{destination}'")
+    resp = requests.get(url, stream=True)
+    if resp.status_code != 200:
+        raise RuntimeError(f"Error downloading {url}")
+    with open(destination, "wb") as file:
+        file.write(resp.raw.read())
+
+
+def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
+    print(f"Extracting '{archive_path}' to '{destination}'")
+    with zipfile.ZipFile(archive_path) as file_in:
+        file_in.extractall(destination)
+
+
+def update_metadata(metadata_path: FileDescriptorOrPath, new_extra_description: str) -> None:
+    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(metadata, file)
+    print(f"Updated {metadata_path}")
+
+
+def run_protoc(
+    proto_paths: Iterable[StrPath], proto_globs: Iterable[str], stubs_folder: StrPath, 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(protoc_version)
+    protoc_args = [
+        *[f"--proto_path={proto_path}" for proto_path in proto_paths],
+        "--mypy_out",
+        f"relax_strict_optional_primitives:{stubs_folder}",
+        *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
diff --git a/scripts/protobuf/s2clientprotocol.py b/scripts/protobuf/s2clientprotocol.py
new file mode 100644
index 000000000000..dd5dfe868880
--- /dev/null
+++ b/scripts/protobuf/s2clientprotocol.py
@@ -0,0 +1,73 @@
+"""
+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 _helpers import download_file, extract_archive, run_protoc, update_metadata
+from mypy_protobuf.main import __version__ as MYPY_PROTOBUF_VERSION  # pyright: ignore[reportMissingTypeStubs]
+
+# Whenever you update S2CLIENT_PROTO_VERSION here, version should be updated
+# in stubs/s2clientprotocol/METADATA.toml and vice-versa.
+S2CLIENT_PROTO_VERSION = "c04df4adbe274858a4eb8417175ee32ad02fd609"
+
+REPO_ROOT = Path(__file__).absolute().parent.parent.parent
+S2CLIENT_STUBS_FOLDER = REPO_ROOT / "stubs" / "s2clientprotocol"
+S2CLIENT_PROTO_FILENAME = f"{S2CLIENT_PROTO_VERSION}.zip"
+S2CLIENT_PROTO_URL = f"https://github.com/Blizzard/s2client-proto/archive/{S2CLIENT_PROTO_FILENAME}"
+S2CLIENT_PROTO_DIR = f"s2client-proto-{S2CLIENT_PROTO_VERSION}"
+VERSION_PATTERN = re.compile(r'def game_version\(\):\n    return "(.+?)"')
+
+
+def extract_version(file_path: Path):
+    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 / S2CLIENT_PROTO_FILENAME
+    download_file(S2CLIENT_PROTO_URL, archive_path)
+    extract_archive(archive_path, temp_dir)
+
+    # Remove existing pyi
+    for old_stub in S2CLIENT_STUBS_FOLDER.rglob("*_pb2.pyi"):
+        old_stub.unlink()
+
+    PROTOC_VERSION = run_protoc(
+        proto_paths=[S2CLIENT_PROTO_DIR],
+        proto_globs=[f"{S2CLIENT_PROTO_DIR}/s2clientprotocol/*.proto"],
+        stubs_folder=S2CLIENT_STUBS_FOLDER,
+        cwd=temp_dir,
+    )
+
+    PYTHON_S2CLIENT_PROTO_VERSION = extract_version(temp_dir / S2CLIENT_PROTO_DIR / "s2clientprotocol" / "build.py")
+
+    # Cleanup
+    shutil.rmtree(temp_dir)
+
+    update_metadata(
+        S2CLIENT_STUBS_FOLDER / "METADATA.toml",
+        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_S2CLIENT_PROTO_VERSION}](https://github.com/Blizzard/s2client-proto/tree/{S2CLIENT_PROTO_VERSION})
+""",
+    )
+
+    # Run pre-commit to cleanup the stubs
+    subprocess.run([sys.executable, "-m", "pre_commit", "run", "--files", *S2CLIENT_STUBS_FOLDER.rglob("*_pb2.pyi")])
+
+
+if __name__ == "__main__":
+    main()
diff --git a/stubs/s2clientprotocol/METADATA.toml b/stubs/s2clientprotocol/METADATA.toml
index c59e1b5b1bf3..9e096cbe23aa 100644
--- a/stubs/s2clientprotocol/METADATA.toml
+++ b/stubs/s2clientprotocol/METADATA.toml
@@ -3,4 +3,4 @@
 version = "5.*"
 upstream_repository = "https://github.com/Blizzard/s2client-proto"
 requires = ["types-protobuf"]
-extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 3.6.1 on [s2client-proto 5.0.12.91115.0](https://github.com/Blizzard/s2client-proto/tree/c04df4adbe274858a4eb8417175ee32ad02fd609)"
+extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on [s2client-proto 5.0.12.91115.0](https://github.com/Blizzard/s2client-proto/tree/c04df4adbe274858a4eb8417175ee32ad02fd609)\n"

From cdd58f348ce96183c94577b8ac1c74afde9fcab0 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Thu, 15 Aug 2024 12:04:10 -0400
Subject: [PATCH 2/9] Completed tensorflow script

---
 .../{protobuf => sync_protobuf}/__init__.py   |   0
 .../{protobuf => sync_protobuf}/_helpers.py   |   8 +-
 .../s2clientprotocol.py                       |  41 +++---
 scripts/sync_protobuf/tensorflow.py           | 138 ++++++++++++++++++
 .../sync_s2clientprotocol_protobuf_stubs.sh   |  62 --------
 scripts/sync_tensorflow_protobuf_stubs.sh     | 101 -------------
 stubs/tensorflow/METADATA.toml                |   2 +-
 7 files changed, 165 insertions(+), 187 deletions(-)
 rename scripts/{protobuf => sync_protobuf}/__init__.py (100%)
 rename scripts/{protobuf => sync_protobuf}/_helpers.py (83%)
 rename scripts/{protobuf => sync_protobuf}/s2clientprotocol.py (50%)
 create mode 100644 scripts/sync_protobuf/tensorflow.py
 delete mode 100644 scripts/sync_s2clientprotocol_protobuf_stubs.sh
 delete mode 100755 scripts/sync_tensorflow_protobuf_stubs.sh

diff --git a/scripts/protobuf/__init__.py b/scripts/sync_protobuf/__init__.py
similarity index 100%
rename from scripts/protobuf/__init__.py
rename to scripts/sync_protobuf/__init__.py
diff --git a/scripts/protobuf/_helpers.py b/scripts/sync_protobuf/_helpers.py
similarity index 83%
rename from scripts/protobuf/_helpers.py
rename to scripts/sync_protobuf/_helpers.py
index f7debe9ac9e7..ab2fc4fcde9a 100644
--- a/scripts/protobuf/_helpers.py
+++ b/scripts/sync_protobuf/_helpers.py
@@ -3,6 +3,7 @@
 import subprocess
 import sys
 import zipfile
+from pathlib import Path
 from typing import TYPE_CHECKING, Iterable
 
 import requests
@@ -11,6 +12,8 @@
 if TYPE_CHECKING:
     from _typeshed import FileDescriptorOrPath, StrOrBytesPath, StrPath
 
+REPO_ROOT = Path(__file__).absolute().parent.parent.parent
+
 
 def download_file(url: str | bytes, destination: FileDescriptorOrPath) -> None:
     print(f"Downloading {url!r} to '{destination}'")
@@ -27,7 +30,8 @@ def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
         file_in.extractall(destination)
 
 
-def update_metadata(metadata_path: FileDescriptorOrPath, new_extra_description: str) -> None:
+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
@@ -37,7 +41,7 @@ def update_metadata(metadata_path: FileDescriptorOrPath, new_extra_description:
 
 
 def run_protoc(
-    proto_paths: Iterable[StrPath], proto_globs: Iterable[str], stubs_folder: StrPath, cwd: StrOrBytesPath | None = None
+    proto_paths: Iterable[StrPath], stubs_folder: StrPath, proto_globs: Iterable[str], cwd: StrOrBytesPath | None = None
 ) -> str:
     """TODO: Describe parameters and return"""
     protoc_version = (
diff --git a/scripts/protobuf/s2clientprotocol.py b/scripts/sync_protobuf/s2clientprotocol.py
similarity index 50%
rename from scripts/protobuf/s2clientprotocol.py
rename to scripts/sync_protobuf/s2clientprotocol.py
index dd5dfe868880..29a5bc998e5d 100644
--- a/scripts/protobuf/s2clientprotocol.py
+++ b/scripts/sync_protobuf/s2clientprotocol.py
@@ -12,22 +12,22 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import download_file, extract_archive, run_protoc, update_metadata
+from _helpers import REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
 from mypy_protobuf.main import __version__ as MYPY_PROTOBUF_VERSION  # pyright: ignore[reportMissingTypeStubs]
 
-# Whenever you update S2CLIENT_PROTO_VERSION here, version should be updated
+# Whenever you update PACKAGE_VERSION here, version should be updated
 # in stubs/s2clientprotocol/METADATA.toml and vice-versa.
-S2CLIENT_PROTO_VERSION = "c04df4adbe274858a4eb8417175ee32ad02fd609"
+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}"
 
-REPO_ROOT = Path(__file__).absolute().parent.parent.parent
-S2CLIENT_STUBS_FOLDER = REPO_ROOT / "stubs" / "s2clientprotocol"
-S2CLIENT_PROTO_FILENAME = f"{S2CLIENT_PROTO_VERSION}.zip"
-S2CLIENT_PROTO_URL = f"https://github.com/Blizzard/s2client-proto/archive/{S2CLIENT_PROTO_FILENAME}"
-S2CLIENT_PROTO_DIR = f"s2client-proto-{S2CLIENT_PROTO_VERSION}"
 VERSION_PATTERN = re.compile(r'def game_version\(\):\n    return "(.+?)"')
 
 
-def extract_version(file_path: Path):
+def extract_version(file_path: Path) -> str:
     match = re.search(VERSION_PATTERN, file_path.read_text())
     assert match
     return match.group(1)
@@ -36,37 +36,36 @@ def extract_version(file_path: Path):
 def main() -> None:
     temp_dir = Path(tempfile.mkdtemp())
     # Fetch s2clientprotocol (which contains all the .proto files)
-    archive_path = temp_dir / S2CLIENT_PROTO_FILENAME
-    download_file(S2CLIENT_PROTO_URL, archive_path)
+    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 S2CLIENT_STUBS_FOLDER.rglob("*_pb2.pyi"):
+    for old_stub in STUBS_FOLDER.rglob("*_pb2.pyi"):
         old_stub.unlink()
 
     PROTOC_VERSION = run_protoc(
-        proto_paths=[S2CLIENT_PROTO_DIR],
-        proto_globs=[f"{S2CLIENT_PROTO_DIR}/s2clientprotocol/*.proto"],
-        stubs_folder=S2CLIENT_STUBS_FOLDER,
+        proto_paths=[EXTRACTED_PACKAGE_DIR],
+        stubs_folder=STUBS_FOLDER,
+        proto_globs=[f"{EXTRACTED_PACKAGE_DIR}/s2clientprotocol/*.proto"],
         cwd=temp_dir,
     )
 
-    PYTHON_S2CLIENT_PROTO_VERSION = extract_version(temp_dir / S2CLIENT_PROTO_DIR / "s2clientprotocol" / "build.py")
+    PYTHON_S2_CLIENT_PROTO_VERSION = extract_version(temp_dir / EXTRACTED_PACKAGE_DIR / "s2clientprotocol" / "build.py")
 
-    # Cleanup
+    # Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
     shutil.rmtree(temp_dir)
 
     update_metadata(
-        S2CLIENT_STUBS_FOLDER / "METADATA.toml",
+        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_S2CLIENT_PROTO_VERSION}](https://github.com/Blizzard/s2client-proto/tree/{S2CLIENT_PROTO_VERSION})
-""",
+[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", *S2CLIENT_STUBS_FOLDER.rglob("*_pb2.pyi")])
+    subprocess.run([sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")])
 
 
 if __name__ == "__main__":
diff --git a/scripts/sync_protobuf/tensorflow.py b/scripts/sync_protobuf/tensorflow.py
new file mode 100644
index 000000000000..b6f9b669cf9f
--- /dev/null
+++ b/scripts/sync_protobuf/tensorflow.py
@@ -0,0 +1,138 @@
+"""
+Generates the protobuf stubs for the given tensorflow version using mypy-protobuf.
+Generally, new minor versions are a good time to update the stubs.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+from _helpers import REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+from mypy_protobuf.main import __version__ as MYPY_PROTOBUF_VERSION  # pyright: ignore[reportMissingTypeStubs]
+
+# Whenever you update PACKAGE_VERSION here, version should be updated
+# in stubs/tensorflow/METADATA.toml and vice-versa.
+PACKAGE_VERSION = "2.17.0"
+
+STUBS_FOLDER = REPO_ROOT / "stubs" / "tensorflow"
+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}"
+
+PROTOS_TO_REMOVE = (
+    "compiler/xla/autotune_results_pb2.pyi",
+    "compiler/xla/autotuning_pb2.pyi",
+    "compiler/xla/service/buffer_assignment_pb2.pyi",
+    "compiler/xla/service/hlo_execution_profile_data_pb2.pyi",
+    "core/protobuf/autotuning_pb2.pyi",
+    "core/protobuf/conv_autotuning_pb2.pyi",
+    "core/protobuf/critical_section_pb2.pyi",
+    "core/protobuf/eager_service_pb2.pyi",
+    "core/protobuf/master_pb2.pyi",
+    "core/protobuf/master_service_pb2.pyi",
+    "core/protobuf/replay_log_pb2.pyi",
+    "core/protobuf/tpu/compile_metadata_pb2.pyi",
+    "core/protobuf/worker_pb2.pyi",
+    "core/protobuf/worker_service_pb2.pyi",
+    "core/util/example_proto_fast_parsing_test_pb2.pyi",
+)
+"""
+These protos exist in a folder with protos used in python,
+but are not included in the python wheel.
+They are likely only used for other language builds.
+stubtest was used to identify them by looking for ModuleNotFoundError.
+(comment out ".*_pb2.*" from the allowlist)
+"""
+
+TSL_IMPORT_PATTERN = re.compile(r"(\[|\s)tsl\.")
+XLA_IMPORT_PATTERN = re.compile(r"(\[|\s)xla\.")
+
+
+def post_creation() -> None:
+    """Move third-party and fix imports"""
+    # Can't use shutil.move because it can't merge existing directories.
+    print()
+    print(f"Moving '{STUBS_FOLDER}/tsl' to '{STUBS_FOLDER}/tensorflow/tsl'")
+    shutil.copytree(f"{STUBS_FOLDER}/tsl", f"{STUBS_FOLDER}/tensorflow/tsl", dirs_exist_ok=True)
+    shutil.rmtree(f"{STUBS_FOLDER}/tsl")
+
+    print(f"Moving '{STUBS_FOLDER}/xla' to '{STUBS_FOLDER}/tensorflow/compiler/xla'")
+    shutil.copytree(f"{STUBS_FOLDER}/xla", f"{STUBS_FOLDER}/tensorflow/compiler/xla", dirs_exist_ok=True)
+    shutil.rmtree(f"{STUBS_FOLDER}/xla")
+
+    for path in STUBS_FOLDER.rglob("*_pb2.pyi"):
+        print(f"Fixing imports in '{path}'")
+        with open(path) as file:
+            filedata = file.read()
+
+        # Replace the target string
+        filedata = re.sub(TSL_IMPORT_PATTERN, "\\1tensorflow.tsl.", filedata)
+        filedata = re.sub(XLA_IMPORT_PATTERN, "\\1tensorflow.compiler.xla.", filedata)
+
+        # Write the file out again
+        with open(path, "w") as file:
+            file.write(filedata)
+
+    print()
+    for to_remove in PROTOS_TO_REMOVE:
+        file_path = STUBS_FOLDER / "tensorflow" / to_remove
+        os.remove(file_path)
+        print(f"Removed '{file_path}'")
+
+
+def main() -> None:
+    temp_dir = Path(tempfile.mkdtemp())
+    # Fetch tensorflow (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}/third_party/xla/third_party/tsl",
+            f"{EXTRACTED_PACKAGE_DIR}/third_party/xla",
+            f"{EXTRACTED_PACKAGE_DIR}",
+        ],
+        stubs_folder=STUBS_FOLDER,
+        proto_globs=[
+            f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/xla/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/xla/service/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/example/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/framework/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/protobuf/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/protobuf/tpu/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/util/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/tensorflow/python/keras/protobuf/*.proto",
+            f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/third_party/tsl/tsl/protobuf/*.proto",
+        ],
+        cwd=temp_dir,
+    )
+
+    # Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
+    shutil.rmtree(temp_dir)
+
+    post_creation()
+
+    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 `tensorflow=={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()
diff --git a/scripts/sync_s2clientprotocol_protobuf_stubs.sh b/scripts/sync_s2clientprotocol_protobuf_stubs.sh
deleted file mode 100644
index 6926c8a0c2b4..000000000000
--- a/scripts/sync_s2clientprotocol_protobuf_stubs.sh
+++ /dev/null
@@ -1,62 +0,0 @@
-#!/bin/bash
-# Based on scripts/generate_proto_stubs.sh.
-# Generates the protobuf stubs for the given s2clientprotocol version using mypy-protobuf.
-# Generally, new minor versions are a good time to update the stubs.
-
-set -euxo pipefail
-
-# Whenever you update S2CLIENT_PROTO_VERSION here, version should be updated
-# in stubs/s2clientprotocol/METADATA.toml and vice-versa.
-S2CLIENT_PROTO_VERSION=c04df4adbe274858a4eb8417175ee32ad02fd609
-MYPY_PROTOBUF_VERSION=3.6.0
-
-REPO_ROOT="$(realpath "$(dirname "${BASH_SOURCE[0]}")"/..)"
-TMP_DIR="$(mktemp -d)"
-S2CLIENT_PROTO_FILENAME="$S2CLIENT_PROTO_VERSION.zip"
-S2CLIENT_PROTO_URL="https://github.com/Blizzard/s2client-proto/archive/$S2CLIENT_PROTO_FILENAME"
-S2CLIENT_PROTO_DIR="s2client-proto-$S2CLIENT_PROTO_VERSION"
-
-cd "$TMP_DIR"
-echo "Working in $TMP_DIR"
-
-# Fetch s2clientprotocol (which contains all the .proto files)
-wget "$S2CLIENT_PROTO_URL"
-unzip "$S2CLIENT_PROTO_FILENAME"
-
-# Prepare virtualenv
-python3 -m venv .venv
-source .venv/bin/activate
-python3 -m pip install pre-commit mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
-
-# Remove existing pyi
-find "$REPO_ROOT/stubs/s2clientprotocol/" -name "*_pb2.pyi" -delete
-
-# s2client works on very old protoc versions, down to 2.6. So we can use the system's protoc.
-PROTOC_VERSION=$(protoc --version)
-echo $PROTOC_VERSION
-protoc \
-  --proto_path="$S2CLIENT_PROTO_DIR" \
-  --mypy_out "relax_strict_optional_primitives:$REPO_ROOT/stubs/s2clientprotocol" \
-  $S2CLIENT_PROTO_DIR/s2clientprotocol/*.proto \
-
-PYTHON_S2CLIENT_PROTO_VERSION=$(
-  grep -Pzo 'def game_version\(\):\n    return ".+?"' $S2CLIENT_PROTO_DIR/s2clientprotocol/build.py \
-  | tr '\n' ' ' \
-  | cut -d '"' -f 2
-)
-
-# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
-rm -rf "$TMP_DIR"
-# Must be in a git repository to run pre-commit
-cd "$REPO_ROOT"
-
-sed -i "" \
-  "s/extra_description = .*$/extra_description = \"\
-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_S2CLIENT_PROTO_VERSION](https:\/\/github.com\/Blizzard\/s2client-proto\/tree\/$S2CLIENT_PROTO_VERSION)\"/" \
-  stubs/s2clientprotocol/METADATA.toml
-
-# use `|| true` so the script still continues even if a pre-commit hook
-# applies autofixes (which will result in a nonzero exit code)
-pre-commit run --files $(git ls-files -- "stubs/s2clientprotocol/**_pb2.pyi") || true
diff --git a/scripts/sync_tensorflow_protobuf_stubs.sh b/scripts/sync_tensorflow_protobuf_stubs.sh
deleted file mode 100755
index 99064cb1e8de..000000000000
--- a/scripts/sync_tensorflow_protobuf_stubs.sh
+++ /dev/null
@@ -1,101 +0,0 @@
-#!/bin/bash
-# Based on scripts/generate_proto_stubs.sh.
-# Generates the protobuf stubs for the given tensorflow version using mypy-protobuf.
-# Generally, new minor versions are a good time to update the stubs.
-
-set -euxo pipefail
-
-# Whenever you update TENSORFLOW_VERSION here, version should be updated
-# in stubs/tensorflow/METADATA.toml and vice-versa.
-TENSORFLOW_VERSION=2.17.0
-MYPY_PROTOBUF_VERSION=3.6.0
-
-# brew install coreutils wget
-# sudo apt-get install -y unzip
-REPO_ROOT="$(realpath "$(dirname "${BASH_SOURCE[0]}")"/..)"
-TMP_DIR="$(mktemp -d)"
-TENSORFLOW_FILENAME="v$TENSORFLOW_VERSION.zip"
-TENSORFLOW_URL="https://github.com/tensorflow/tensorflow/archive/refs/tags/$TENSORFLOW_FILENAME"
-TENSORFLOW_DIR="tensorflow-$TENSORFLOW_VERSION"
-
-cd "$TMP_DIR"
-echo "Working in $TMP_DIR"
-
-# Fetch tensorflow (which contains all the .proto files)
-wget "$TENSORFLOW_URL"
-unzip "$TENSORFLOW_FILENAME"
-
-# Prepare virtualenv
-python3 -m venv .venv
-source .venv/bin/activate
-python3 -m pip install grpcio-tools pre-commit mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
-
-# Empty target folders or the mv command below will fail
-rm -rf "$REPO_ROOT/stubs/tensorflow/tensorflow/tsl/"
-rm -rf "$REPO_ROOT/stubs/tensorflow/tensorflow/compiler/xla/"
-# Remove existing pyi
-find "$REPO_ROOT/stubs/tensorflow/" -name "*_pb2.pyi" -delete
-
-# Folders here cover the more commonly used protobufs externally and
-# their dependencies. Tensorflow has more protobufs and can be added if requested.
-PROTOC_VERSION=$(python3 -m grpc_tools.protoc --version)
-echo $PROTOC_VERSION
-python3 -m grpc_tools.protoc \
-  --proto_path="$TENSORFLOW_DIR/third_party/xla/third_party/tsl" \
-  --proto_path="$TENSORFLOW_DIR/third_party/xla" \
-  --proto_path="$TENSORFLOW_DIR" \
-  --mypy_out "relax_strict_optional_primitives:$REPO_ROOT/stubs/tensorflow" \
-  $TENSORFLOW_DIR/third_party/xla/xla/*.proto \
-  $TENSORFLOW_DIR/third_party/xla/xla/service/*.proto \
-  $TENSORFLOW_DIR/tensorflow/core/example/*.proto \
-  $TENSORFLOW_DIR/tensorflow/core/framework/*.proto \
-  $TENSORFLOW_DIR/tensorflow/core/protobuf/*.proto \
-  $TENSORFLOW_DIR/tensorflow/core/protobuf/tpu/*.proto \
-  $TENSORFLOW_DIR/tensorflow/core/util/*.proto \
-  $TENSORFLOW_DIR/tensorflow/python/keras/protobuf/*.proto \
-  $TENSORFLOW_DIR/third_party/xla/third_party/tsl/tsl/protobuf/*.proto \
-
-# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
-rm -rf "$TMP_DIR"
-
-# Must be in a git repository to run pre-commit
-cd "$REPO_ROOT"
-
-# Move third-party and fix imports
-mv stubs/tensorflow/tsl/ stubs/tensorflow/tensorflow/
-find stubs/tensorflow/ -name '*_pb2.pyi' | xargs sed -i "" -r "s/(\[|\s)tsl\./\1tensorflow\.tsl\./"
-mv stubs/tensorflow/xla/ stubs/tensorflow/tensorflow/compiler/
-find stubs/tensorflow/ -name '*_pb2.pyi' | xargs sed -i "" -r "s/(\[|\s)xla\./\1tensorflow\.compiler\.xla\./"
-
-# These protos exist in a folder with protos used in python,
-# but are not included in the python wheel.
-# They are likely only used for other language builds.
-# stubtest was used to identify them by looking for ModuleNotFoundError.
-# (comment out ".*_pb2.*" from the allowlist)
-rm -r \
-  stubs/tensorflow/tensorflow/compiler/xla/autotune_results_pb2.pyi \
-  stubs/tensorflow/tensorflow/compiler/xla/autotuning_pb2.pyi \
-  stubs/tensorflow/tensorflow/compiler/xla/service/buffer_assignment_pb2.pyi \
-  stubs/tensorflow/tensorflow/compiler/xla/service/hlo_execution_profile_data_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/autotuning_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/conv_autotuning_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/critical_section_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/eager_service_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/master_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/master_service_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/replay_log_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/tpu/compile_metadata_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/worker_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/protobuf/worker_service_pb2.pyi \
-  stubs/tensorflow/tensorflow/core/util/example_proto_fast_parsing_test_pb2.pyi \
-
-sed -i "" \
-  "s/extra_description = .*$/extra_description = \"\
-Partially generated using [mypy-protobuf==$MYPY_PROTOBUF_VERSION](https:\/\/github.com\/nipunn1313\/mypy-protobuf\/tree\/v$MYPY_PROTOBUF_VERSION) \
-and $PROTOC_VERSION \
-on tensorflow==$TENSORFLOW_VERSION .\"/" \
-  stubs/tensorflow/METADATA.toml
-
-# use `|| true` so the script still continues even if a pre-commit hook
-# applies autofixes (which will result in a nonzero exit code)
-pre-commit run --files $(git ls-files -- "stubs/tensorflow/**_pb2.pyi") || true
diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml
index 0cf0913eda0a..c8dd6cadc04d 100644
--- a/stubs/tensorflow/METADATA.toml
+++ b/stubs/tensorflow/METADATA.toml
@@ -4,7 +4,7 @@ version = "2.17.*"
 upstream_repository = "https://github.com/tensorflow/tensorflow"
 # requires a version of numpy with a `py.typed` file
 requires = ["numpy>=1.20", "types-protobuf", "types-requests"]
-extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on tensorflow==2.17.0 ."
+extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on `tensorflow==2.17.0`.)"
 partial_stub = true
 
 [tool.stubtest]

From 9f33b0edd3ab0ffae18b92e360059aa0be0d2328 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Thu, 15 Aug 2024 12:28:30 -0400
Subject: [PATCH 3/9] Fix CI tests

---
 requirements-tests.txt                    |  1 -
 scripts/sync_protobuf/__init__.py         |  0
 scripts/sync_protobuf/_helpers.py         | 24 ++++++++++++++---------
 scripts/sync_protobuf/s2clientprotocol.py | 11 +++++------
 scripts/sync_protobuf/tensorflow.py       | 15 +++++++-------
 stubs/s2clientprotocol/METADATA.toml      |  2 +-
 6 files changed, 28 insertions(+), 25 deletions(-)
 delete mode 100644 scripts/sync_protobuf/__init__.py

diff --git a/requirements-tests.txt b/requirements-tests.txt
index 99e05ecfe5a8..a8cf6bbc5ec9 100644
--- a/requirements-tests.txt
+++ b/requirements-tests.txt
@@ -18,7 +18,6 @@ mypy-protobuf==3.6.0
 packaging==24.1
 pathspec>=0.11.1
 pre-commit
-requests
 stubdefaulter==0.1.0
 termcolor>=2.3
 tomli==2.0.1
diff --git a/scripts/sync_protobuf/__init__.py b/scripts/sync_protobuf/__init__.py
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/scripts/sync_protobuf/_helpers.py b/scripts/sync_protobuf/_helpers.py
index ab2fc4fcde9a..4adf889b7a91 100644
--- a/scripts/sync_protobuf/_helpers.py
+++ b/scripts/sync_protobuf/_helpers.py
@@ -3,25 +3,30 @@
 import subprocess
 import sys
 import zipfile
+from http.client import HTTPResponse
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterable
+from urllib.request import urlopen
 
-import requests
 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 FileDescriptorOrPath, StrOrBytesPath, StrPath
 
 REPO_ROOT = Path(__file__).absolute().parent.parent.parent
+MYPY_PROTOBUF_VERSION = mypy_protobuf__version__
 
 
-def download_file(url: str | bytes, destination: FileDescriptorOrPath) -> None:
+def download_file(url: str, destination: FileDescriptorOrPath) -> None:
     print(f"Downloading {url!r} to '{destination}'")
-    resp = requests.get(url, stream=True)
-    if resp.status_code != 200:
+    resp: HTTPResponse = urlopen(url)
+    if resp.getcode() != 200:
         raise RuntimeError(f"Error downloading {url}")
     with open(destination, "wb") as file:
-        file.write(resp.raw.read())
+        file.write(resp.read())
 
 
 def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
@@ -36,12 +41,13 @@ def update_metadata(metadata_folder: StrPath, new_extra_description: str) -> Non
         metadata = tomlkit.load(file)
     metadata["extra_description"] = new_extra_description
     with open(metadata_path, "w") as file:
-        tomlkit.dump(metadata, 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], stubs_folder: StrPath, proto_globs: Iterable[str], cwd: StrOrBytesPath | None = None
+    proto_paths: Iterable[StrPath], mypy_out: StrPath, proto_globs: Iterable[str], cwd: StrOrBytesPath | None = None
 ) -> str:
     """TODO: Describe parameters and return"""
     protoc_version = (
@@ -51,9 +57,9 @@ def run_protoc(
     protoc_args = [
         *[f"--proto_path={proto_path}" for proto_path in proto_paths],
         "--mypy_out",
-        f"relax_strict_optional_primitives:{stubs_folder}",
+        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)
+    subprocess.run((sys.executable, "-m", "grpc_tools.protoc", *protoc_args), cwd=cwd, check=True)
     return protoc_version
diff --git a/scripts/sync_protobuf/s2clientprotocol.py b/scripts/sync_protobuf/s2clientprotocol.py
index 29a5bc998e5d..5e8d8750204f 100644
--- a/scripts/sync_protobuf/s2clientprotocol.py
+++ b/scripts/sync_protobuf/s2clientprotocol.py
@@ -12,8 +12,7 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
-from mypy_protobuf.main import __version__ as MYPY_PROTOBUF_VERSION  # pyright: ignore[reportMissingTypeStubs]
+from _helpers 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.
@@ -45,9 +44,9 @@ def main() -> None:
         old_stub.unlink()
 
     PROTOC_VERSION = run_protoc(
-        proto_paths=[EXTRACTED_PACKAGE_DIR],
-        stubs_folder=STUBS_FOLDER,
-        proto_globs=[f"{EXTRACTED_PACKAGE_DIR}/s2clientprotocol/*.proto"],
+        proto_paths=(EXTRACTED_PACKAGE_DIR,),
+        mypy_out=STUBS_FOLDER,
+        proto_globs=(f"{EXTRACTED_PACKAGE_DIR}/s2clientprotocol/*.proto",),
         cwd=temp_dir,
     )
 
@@ -65,7 +64,7 @@ def main() -> None:
     )
 
     # Run pre-commit to cleanup the stubs
-    subprocess.run([sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")])
+    subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
 
 
 if __name__ == "__main__":
diff --git a/scripts/sync_protobuf/tensorflow.py b/scripts/sync_protobuf/tensorflow.py
index b6f9b669cf9f..28f3832bf9ec 100644
--- a/scripts/sync_protobuf/tensorflow.py
+++ b/scripts/sync_protobuf/tensorflow.py
@@ -13,8 +13,7 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
-from mypy_protobuf.main import __version__ as MYPY_PROTOBUF_VERSION  # pyright: ignore[reportMissingTypeStubs]
+from _helpers 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/tensorflow/METADATA.toml and vice-versa.
@@ -98,13 +97,13 @@ def main() -> None:
         old_stub.unlink()
 
     PROTOC_VERSION = run_protoc(
-        proto_paths=[
+        proto_paths=(
             f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/third_party/tsl",
             f"{EXTRACTED_PACKAGE_DIR}/third_party/xla",
             f"{EXTRACTED_PACKAGE_DIR}",
-        ],
-        stubs_folder=STUBS_FOLDER,
-        proto_globs=[
+        ),
+        mypy_out=STUBS_FOLDER,
+        proto_globs=(
             f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/xla/*.proto",
             f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/xla/service/*.proto",
             f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/example/*.proto",
@@ -114,7 +113,7 @@ def main() -> None:
             f"{EXTRACTED_PACKAGE_DIR}/tensorflow/core/util/*.proto",
             f"{EXTRACTED_PACKAGE_DIR}/tensorflow/python/keras/protobuf/*.proto",
             f"{EXTRACTED_PACKAGE_DIR}/third_party/xla/third_party/tsl/tsl/protobuf/*.proto",
-        ],
+        ),
         cwd=temp_dir,
     )
 
@@ -131,7 +130,7 @@ def main() -> None:
     )
 
     # Run pre-commit to cleanup the stubs
-    subprocess.run([sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")])
+    subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
 
 
 if __name__ == "__main__":
diff --git a/stubs/s2clientprotocol/METADATA.toml b/stubs/s2clientprotocol/METADATA.toml
index 9e096cbe23aa..97e85b24535b 100644
--- a/stubs/s2clientprotocol/METADATA.toml
+++ b/stubs/s2clientprotocol/METADATA.toml
@@ -3,4 +3,4 @@
 version = "5.*"
 upstream_repository = "https://github.com/Blizzard/s2client-proto"
 requires = ["types-protobuf"]
-extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on [s2client-proto 5.0.12.91115.0](https://github.com/Blizzard/s2client-proto/tree/c04df4adbe274858a4eb8417175ee32ad02fd609)\n"
+extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on [s2client-proto 5.0.12.91115.0](https://github.com/Blizzard/s2client-proto/tree/c04df4adbe274858a4eb8417175ee32ad02fd609)."

From c1ba6e5e30f9a4c925af190d217125f2ca44ac69 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Fri, 16 Aug 2024 15:16:33 -0400
Subject: [PATCH 4/9] Finished google protobuf

---
 scripts/sync_protobuf/_helpers.py         | 11 +--
 scripts/sync_protobuf/google_protobuf.py  | 91 +++++++++++++++++++++++
 scripts/sync_protobuf/s2clientprotocol.py |  5 +-
 stubs/protobuf/METADATA.toml              |  6 +-
 stubs/s2clientprotocol/METADATA.toml      |  4 +-
 stubs/tensorflow/METADATA.toml            |  4 +-
 6 files changed, 107 insertions(+), 14 deletions(-)
 create mode 100644 scripts/sync_protobuf/google_protobuf.py

diff --git a/scripts/sync_protobuf/_helpers.py b/scripts/sync_protobuf/_helpers.py
index 4adf889b7a91..a5d8be465078 100644
--- a/scripts/sync_protobuf/_helpers.py
+++ b/scripts/sync_protobuf/_helpers.py
@@ -2,11 +2,11 @@
 
 import subprocess
 import sys
-import zipfile
 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]
@@ -14,14 +14,14 @@
 )
 
 if TYPE_CHECKING:
-    from _typeshed import FileDescriptorOrPath, StrOrBytesPath, StrPath
+    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: FileDescriptorOrPath) -> None:
-    print(f"Downloading {url!r} to '{destination}'")
+def download_file(url: str, destination: StrPath) -> None:
+    print(f"Downloading '{url}' to '{destination}'")
     resp: HTTPResponse = urlopen(url)
     if resp.getcode() != 200:
         raise RuntimeError(f"Error downloading {url}")
@@ -31,7 +31,7 @@ def download_file(url: str, destination: FileDescriptorOrPath) -> None:
 
 def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
     print(f"Extracting '{archive_path}' to '{destination}'")
-    with zipfile.ZipFile(archive_path) as file_in:
+    with ZipFile(archive_path) as file_in:
         file_in.extractall(destination)
 
 
@@ -53,6 +53,7 @@ def run_protoc(
     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],
diff --git a/scripts/sync_protobuf/google_protobuf.py b/scripts/sync_protobuf/google_protobuf.py
new file mode 100644
index 000000000000..794aa430c3b9
--- /dev/null
+++ b/scripts/sync_protobuf/google_protobuf.py
@@ -0,0 +1,91 @@
+"""
+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 _helpers 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, dict[str, dict[str, str]]] = json.load(file)
+    # The root key will be the protobuf source code version
+    return next(iter(data.values()))["languages"]["python"]
+
+
+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.readlines()))
+        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()
diff --git a/scripts/sync_protobuf/s2clientprotocol.py b/scripts/sync_protobuf/s2clientprotocol.py
index 5e8d8750204f..22b469d24004 100644
--- a/scripts/sync_protobuf/s2clientprotocol.py
+++ b/scripts/sync_protobuf/s2clientprotocol.py
@@ -26,7 +26,8 @@
 VERSION_PATTERN = re.compile(r'def game_version\(\):\n    return "(.+?)"')
 
 
-def extract_version(file_path: Path) -> str:
+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)
@@ -50,7 +51,7 @@ def main() -> None:
         cwd=temp_dir,
     )
 
-    PYTHON_S2_CLIENT_PROTO_VERSION = extract_version(temp_dir / EXTRACTED_PACKAGE_DIR / "s2clientprotocol" / "build.py")
+    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)
diff --git a/stubs/protobuf/METADATA.toml b/stubs/protobuf/METADATA.toml
index 9e402da86425..f08d51e7fcab 100644
--- a/stubs/protobuf/METADATA.toml
+++ b/stubs/protobuf/METADATA.toml
@@ -1,8 +1,8 @@
-# Whenever you update version here, PROTOBUF_VERSION should be updated
-# in scripts/generate_proto_stubs.sh and vice-versa.
+# Whenever you update version here, PACKAGE_VERSION should be updated
+# in scripts/sync_proto/google_protobuf.py and vice-versa.
 version = "5.27.*"
 upstream_repository = "https://github.com/protocolbuffers/protobuf"
-extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on [protobuf v27.1](https://github.com/protocolbuffers/protobuf/releases/tag/v27.1) (python protobuf==5.27.1)."
+extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on [protobuf v27.1](https://github.com/protocolbuffers/protobuf/releases/tag/v27.1) (python `protobuf==5.27.1`)."
 partial_stub = true
 
 [tool.stubtest]
diff --git a/stubs/s2clientprotocol/METADATA.toml b/stubs/s2clientprotocol/METADATA.toml
index 97e85b24535b..ed86f22032af 100644
--- a/stubs/s2clientprotocol/METADATA.toml
+++ b/stubs/s2clientprotocol/METADATA.toml
@@ -1,5 +1,5 @@
-# Whenever you update version here, S2CLIENT_PROTO_VERSION should be updated
-# in scripts/sync_s2clientprotocol_protobuf_stubs.sh and vice-versa.
+# Whenever you update version here, PACKAGE_VERSION should be updated
+# in scripts/sync_proto/s2clientprotocol.py and vice-versa.
 version = "5.*"
 upstream_repository = "https://github.com/Blizzard/s2client-proto"
 requires = ["types-protobuf"]
diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml
index c8dd6cadc04d..aba8dcab4f0d 100644
--- a/stubs/tensorflow/METADATA.toml
+++ b/stubs/tensorflow/METADATA.toml
@@ -1,5 +1,5 @@
-# Whenever you update version here, TENSORFLOW_VERSION should be updated
-# in scripts/sync_tensorflow_protobuf_stubs.sh and vice-versa.
+# Whenever you update version here, PACKAGE_VERSION should be updated
+# in scripts/sync_proto/tensorflow.py and vice-versa.
 version = "2.17.*"
 upstream_repository = "https://github.com/tensorflow/tensorflow"
 # requires a version of numpy with a `py.typed` file

From 5f2ec266d06d720d0f84e79efc31959c7461e966 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Fri, 16 Aug 2024 15:26:00 -0400
Subject: [PATCH 5/9] Forgot to delete old script

---
 scripts/generate_proto_stubs.sh | 76 ---------------------------------
 1 file changed, 76 deletions(-)
 delete mode 100755 scripts/generate_proto_stubs.sh

diff --git a/scripts/generate_proto_stubs.sh b/scripts/generate_proto_stubs.sh
deleted file mode 100755
index e6f69ae9e7d7..000000000000
--- a/scripts/generate_proto_stubs.sh
+++ /dev/null
@@ -1,76 +0,0 @@
-#!/usr/bin/env bash
-# Some of the proto .pyi stubs in stubs/protobuf/
-# are autogenerated using the mypy-protobuf project on the
-# latest `.proto` files shipped with protoc.
-
-set -ex -o pipefail
-
-# When run, this script will autogenerate the _pb2.pyi stubs to
-# stubs/protobuf. It should be run any time there's
-# a meaningful update to either PROTOBUF_VERSION or MYPY_PROTOBUF_VERSION,
-# followed by committing the changes to typeshed
-#
-# Whenever you update PROTOBUF_VERSION here, version should be updated
-# in stubs/protobuf/METADATA.toml and vice-versa.
-PROTOBUF_VERSION=27.1
-MYPY_PROTOBUF_VERSION=3.6.0
-
-# brew install coreutils wget
-# sudo apt-get install -y unzip
-REPO_ROOT="$(realpath "$(dirname "${BASH_SOURCE[0]}")"/..)"
-TMP_DIR="$(mktemp -d)"
-PYTHON_PROTOBUF_FILENAME="protobuf-$PROTOBUF_VERSION.zip"
-PYTHON_PROTOBUF_URL="https://github.com/protocolbuffers/protobuf/releases/download/v$PROTOBUF_VERSION/$PYTHON_PROTOBUF_FILENAME"
-PYTHON_PROTOBUF_DIR="protobuf-$PROTOBUF_VERSION"
-
-cd "$TMP_DIR"
-echo "Working in $TMP_DIR"
-
-# Fetch protoc-python (which contains all the .proto files)
-wget "$PYTHON_PROTOBUF_URL"
-unzip "$PYTHON_PROTOBUF_FILENAME"
-
-# Prepare virtualenv
-python3 -m venv .venv
-source .venv/bin/activate
-python3 -m pip install grpcio-tools pre-commit mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
-
-# Remove existing pyi
-find "$REPO_ROOT/stubs/protobuf/" -name '*_pb2.pyi' -delete
-
-# 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
-PROTO_FILES=$(grep '"//:.*_proto"' $PYTHON_PROTOBUF_DIR/python/dist/BUILD.bazel | \
-    cut -d\" -f2 | \
-    sed "s://\::$PYTHON_PROTOBUF_DIR/src/google/protobuf/:" | \
-    sed "s:_proto:.proto:" | \
-    sed "s:compiler_:compiler/:" \
-)
-
-# And regenerate!
-PROTOC_VERSION=$(python3 -m grpc_tools.protoc --version)
-echo $PROTOC_VERSION
-python3 -m grpc_tools.protoc \
-  --proto_path="$PYTHON_PROTOBUF_DIR/src" \
-  --mypy_out="relax_strict_optional_primitives:$REPO_ROOT/stubs/protobuf" \
-  $PROTO_FILES
-
-PYTHON_PROTOBUF_VERSION=$(jq -r '.[] | .languages.python' "$PYTHON_PROTOBUF_DIR/version.json")
-
-# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
-rm -rf "$TMP_DIR"
-# Must be in a git repository to run pre-commit
-cd "$REPO_ROOT"
-
-sed -i "" \
-  "s/extra_description = .*$/extra_description = \"\
-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$PROTOBUF_VERSION](https:\/\/github.com\/protocolbuffers\/protobuf\/releases\/tag\/v$PROTOBUF_VERSION) \
-(python protobuf==$PYTHON_PROTOBUF_VERSION).\"/" \
-  stubs/protobuf/METADATA.toml
-
-# use `|| true` so the script still continues even if a pre-commit hook
-# applies autofixes (which will result in a nonzero exit code)
-pre-commit run --files $(git ls-files -- "stubs/protobuf/**_pb2.pyi") || true

From 893a22320319e957ed543180840b614d6255aa9e Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Fri, 16 Aug 2024 15:35:00 -0400
Subject: [PATCH 6/9] Don't install grpcio-tools on Python 3.13

---
 requirements-tests.txt | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/requirements-tests.txt b/requirements-tests.txt
index a8cf6bbc5ec9..9e0817e9a817 100644
--- a/requirements-tests.txt
+++ b/requirements-tests.txt
@@ -13,7 +13,9 @@ ruff==0.5.4              # must match .pre-commit-config.yaml
 
 # Libraries used by our various scripts.
 aiohttp==3.10.2
-grpcio-tools
+# 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

From 75a746c0547627375be8f5d426f89b63a147da70 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Mon, 19 Aug 2024 11:50:06 -0400
Subject: [PATCH 7/9] Rename _helpers to _utils

---
 scripts/sync_protobuf/{_helpers.py => _utils.py} | 0
 scripts/sync_protobuf/google_protobuf.py         | 2 +-
 scripts/sync_protobuf/s2clientprotocol.py        | 2 +-
 scripts/sync_protobuf/tensorflow.py              | 4 ++--
 4 files changed, 4 insertions(+), 4 deletions(-)
 rename scripts/sync_protobuf/{_helpers.py => _utils.py} (100%)

diff --git a/scripts/sync_protobuf/_helpers.py b/scripts/sync_protobuf/_utils.py
similarity index 100%
rename from scripts/sync_protobuf/_helpers.py
rename to scripts/sync_protobuf/_utils.py
diff --git a/scripts/sync_protobuf/google_protobuf.py b/scripts/sync_protobuf/google_protobuf.py
index 794aa430c3b9..53b79bf0d54e 100644
--- a/scripts/sync_protobuf/google_protobuf.py
+++ b/scripts/sync_protobuf/google_protobuf.py
@@ -13,7 +13,7 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+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.
diff --git a/scripts/sync_protobuf/s2clientprotocol.py b/scripts/sync_protobuf/s2clientprotocol.py
index 22b469d24004..93b9c87d544c 100644
--- a/scripts/sync_protobuf/s2clientprotocol.py
+++ b/scripts/sync_protobuf/s2clientprotocol.py
@@ -12,7 +12,7 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+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.
diff --git a/scripts/sync_protobuf/tensorflow.py b/scripts/sync_protobuf/tensorflow.py
index 28f3832bf9ec..2d6005674b3f 100644
--- a/scripts/sync_protobuf/tensorflow.py
+++ b/scripts/sync_protobuf/tensorflow.py
@@ -13,7 +13,7 @@
 import tempfile
 from pathlib import Path
 
-from _helpers import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
+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/tensorflow/METADATA.toml and vice-versa.
@@ -126,7 +126,7 @@ def main() -> None:
         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 `tensorflow=={PACKAGE_VERSION}`.)""",
+and {PROTOC_VERSION} on `tensorflow=={PACKAGE_VERSION}`.""",
     )
 
     # Run pre-commit to cleanup the stubs

From ae26a36221444e04f1bb70e2be9cb4ad5a221530 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Wed, 18 Sep 2024 19:38:20 -0400
Subject: [PATCH 8/9] Apply suggestions from code review

Co-authored-by: Sebastian Rittau <srittau@rittau.biz>
---
 scripts/sync_protobuf/_utils.py          | 11 ++++++-----
 scripts/sync_protobuf/google_protobuf.py |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/scripts/sync_protobuf/_utils.py b/scripts/sync_protobuf/_utils.py
index a5d8be465078..cd97ef67a450 100644
--- a/scripts/sync_protobuf/_utils.py
+++ b/scripts/sync_protobuf/_utils.py
@@ -22,11 +22,12 @@
 
 def download_file(url: str, destination: StrPath) -> None:
     print(f"Downloading '{url}' to '{destination}'")
-    resp: HTTPResponse = urlopen(url)
-    if resp.getcode() != 200:
-        raise RuntimeError(f"Error downloading {url}")
-    with open(destination, "wb") as file:
-        file.write(resp.read())
+    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:
diff --git a/scripts/sync_protobuf/google_protobuf.py b/scripts/sync_protobuf/google_protobuf.py
index 53b79bf0d54e..87c049e85ff9 100644
--- a/scripts/sync_protobuf/google_protobuf.py
+++ b/scripts/sync_protobuf/google_protobuf.py
@@ -43,7 +43,7 @@ def extract_proto_file_paths(temp_dir: Path) -> list[str]:
     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.readlines()))
+        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

From 598d0a775631f3c6c6c48a2a1aa15ea22707c0a4 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Wed, 18 Sep 2024 19:39:42 -0400
Subject: [PATCH 9/9] assert google_protobuf.extract_python_version is a str

---
 scripts/sync_protobuf/google_protobuf.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/scripts/sync_protobuf/google_protobuf.py b/scripts/sync_protobuf/google_protobuf.py
index 87c049e85ff9..ba9109008d54 100644
--- a/scripts/sync_protobuf/google_protobuf.py
+++ b/scripts/sync_protobuf/google_protobuf.py
@@ -12,6 +12,7 @@
 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
 
@@ -31,9 +32,11 @@
 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, dict[str, dict[str, str]]] = json.load(file)
+        data: dict[str, Any] = json.load(file)
     # The root key will be the protobuf source code version
-    return next(iter(data.values()))["languages"]["python"]
+    version = next(iter(data.values()))["languages"]["python"]
+    assert isinstance(version, str)
+    return version
 
 
 def extract_proto_file_paths(temp_dir: Path) -> list[str]: