Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c025e37

Browse files
AvasamAvasam
authored
Rewrote protobuf generation scripts in Python (#12527)
1 parent 0689736 commit c025e37

11 files changed

Lines changed: 383 additions & 248 deletions

‎requirements-tests.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ ruff==0.5.4 # must match .pre-commit-config.yaml
1414

1515
# Libraries used by our various scripts.
1616
aiohttp==3.10.2
17+
# grpc install only fails on Windows, but let's avoid building sdist on other platforms
18+
# https://github.com/grpc/grpc/issues/36201
19+
grpcio-tools; python_version < "3.13"
20+
mypy-protobuf==3.6.0
1721
packaging==24.1
1822
pathspec>=0.11.1
1923
pre-commit

‎scripts/generate_proto_stubs.sh‎

Lines changed: 0 additions & 76 deletions
This file was deleted.

‎scripts/sync_protobuf/_utils.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
from __future__ import annotations
2+
3+
import subprocess
4+
import sys
5+
from http.client import HTTPResponse
6+
from pathlib import Path
7+
from typing import TYPE_CHECKING, Iterable
8+
from urllib.request import urlopen
9+
from zipfile import ZipFile
10+
11+
import tomlkit
12+
from mypy_protobuf.main import ( # type: ignore[import-untyped] # pyright: ignore[reportMissingTypeStubs]
13+
__version__ as mypy_protobuf__version__,
14+
)
15+
16+
if TYPE_CHECKING:
17+
from _typeshed import StrOrBytesPath, StrPath
18+
19+
REPO_ROOT = Path(__file__).absolute().parent.parent.parent
20+
MYPY_PROTOBUF_VERSION = mypy_protobuf__version__
21+
22+
23+
def download_file(url: str, destination: StrPath) -> None:
24+
print(f"Downloading '{url}' to '{destination}'")
25+
resp: HTTPResponse
26+
with urlopen(url) as resp:
27+
if resp.getcode() != 200:
28+
raise RuntimeError(f"Error downloading {url}")
29+
with open(destination, "wb") as file:
30+
file.write(resp.read())
31+
32+
33+
def extract_archive(archive_path: StrPath, destination: StrPath) -> None:
34+
print(f"Extracting '{archive_path}' to '{destination}'")
35+
with ZipFile(archive_path) as file_in:
36+
file_in.extractall(destination)
37+
38+
39+
def update_metadata(metadata_folder: StrPath, new_extra_description: str) -> None:
40+
metadata_path = Path(metadata_folder) / "METADATA.toml"
41+
with open(metadata_path) as file:
42+
metadata = tomlkit.load(file)
43+
metadata["extra_description"] = new_extra_description
44+
with open(metadata_path, "w") as file:
45+
# tomlkit.dump has partially unknown IO type
46+
tomlkit.dump(metadata, file) # pyright: ignore[reportUnknownMemberType]
47+
print(f"Updated {metadata_path}")
48+
49+
50+
def run_protoc(
51+
proto_paths: Iterable[StrPath], mypy_out: StrPath, proto_globs: Iterable[str], cwd: StrOrBytesPath | None = None
52+
) -> str:
53+
"""TODO: Describe parameters and return"""
54+
protoc_version = (
55+
subprocess.run([sys.executable, "-m", "grpc_tools.protoc", "--version"], capture_output=True).stdout.decode().strip()
56+
)
57+
print()
58+
print(protoc_version)
59+
protoc_args = [
60+
*[f"--proto_path={proto_path}" for proto_path in proto_paths],
61+
"--mypy_out",
62+
f"relax_strict_optional_primitives:{mypy_out}",
63+
*proto_globs,
64+
]
65+
print("Running: protoc\n " + "\n ".join(protoc_args) + "\n")
66+
subprocess.run((sys.executable, "-m", "grpc_tools.protoc", *protoc_args), cwd=cwd, check=True)
67+
return protoc_version

‎scripts/sync_protobuf/google_protobuf.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""
2+
Generates the protobuf stubs for the given protobuf version using mypy-protobuf.
3+
Generally, new minor versions are a good time to update the stubs.
4+
"""
5+
6+
from __future__ import annotations
7+
8+
import json
9+
import re
10+
import shutil
11+
import subprocess
12+
import sys
13+
import tempfile
14+
from pathlib import Path
15+
from typing import Any
16+
17+
from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
18+
19+
# Whenever you update PACKAGE_VERSION here, version should be updated
20+
# in stubs/protobuf/METADATA.toml and vice-versa.
21+
PACKAGE_VERSION = "27.1"
22+
23+
STUBS_FOLDER = REPO_ROOT / "stubs" / "protobuf"
24+
ARCHIVE_FILENAME = f"protobuf-{PACKAGE_VERSION}.zip"
25+
ARCHIVE_URL = f"https://github.com/protocolbuffers/protobuf/releases/download/v{PACKAGE_VERSION}/{ARCHIVE_FILENAME}"
26+
EXTRACTED_PACKAGE_DIR = f"protobuf-{PACKAGE_VERSION}"
27+
28+
VERSION_PATTERN = re.compile(r'def game_version\(\):\n return "(.+?)"')
29+
PROTO_FILE_PATTERN = re.compile(r'"//:(.*)_proto"')
30+
31+
32+
def extract_python_version(file_path: Path) -> str:
33+
"""Extract the Python version from https://github.com/protocolbuffers/protobuf/blob/main/version.json"""
34+
with open(file_path) as file:
35+
data: dict[str, Any] = json.load(file)
36+
# The root key will be the protobuf source code version
37+
version = next(iter(data.values()))["languages"]["python"]
38+
assert isinstance(version, str)
39+
return version
40+
41+
42+
def extract_proto_file_paths(temp_dir: Path) -> list[str]:
43+
"""
44+
Roughly reproduce the subset of .proto files on the public interface
45+
as described in py_proto_library calls in
46+
https://github.com/protocolbuffers/protobuf/blob/main/python/dist/BUILD.bazel
47+
"""
48+
with open(temp_dir / EXTRACTED_PACKAGE_DIR / "python" / "dist" / "BUILD.bazel") as file:
49+
matched_lines = filter(None, (re.search(PROTO_FILE_PATTERN, line) for line in file))
50+
proto_files = [
51+
EXTRACTED_PACKAGE_DIR + "/src/google/protobuf/" + match.group(1).replace("compiler_", "compiler/") + ".proto"
52+
for match in matched_lines
53+
]
54+
return proto_files
55+
56+
57+
def main() -> None:
58+
temp_dir = Path(tempfile.mkdtemp())
59+
# Fetch s2clientprotocol (which contains all the .proto files)
60+
archive_path = temp_dir / ARCHIVE_FILENAME
61+
download_file(ARCHIVE_URL, archive_path)
62+
extract_archive(archive_path, temp_dir)
63+
64+
# Remove existing pyi
65+
for old_stub in STUBS_FOLDER.rglob("*_pb2.pyi"):
66+
old_stub.unlink()
67+
68+
PROTOC_VERSION = run_protoc(
69+
proto_paths=(f"{EXTRACTED_PACKAGE_DIR}/src",),
70+
mypy_out=STUBS_FOLDER,
71+
proto_globs=extract_proto_file_paths(temp_dir),
72+
cwd=temp_dir,
73+
)
74+
75+
PYTHON_PROTOBUF_VERSION = extract_python_version(temp_dir / EXTRACTED_PACKAGE_DIR / "version.json")
76+
77+
# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
78+
shutil.rmtree(temp_dir)
79+
80+
update_metadata(
81+
STUBS_FOLDER,
82+
f"""Partially generated using \
83+
[mypy-protobuf=={MYPY_PROTOBUF_VERSION}](https://github.com/nipunn1313/mypy-protobuf/tree/v{MYPY_PROTOBUF_VERSION}) \
84+
and {PROTOC_VERSION} on \
85+
[protobuf v{PACKAGE_VERSION}](https://github.com/protocolbuffers/protobuf/releases/tag/v{PACKAGE_VERSION}) \
86+
(python `protobuf=={PYTHON_PROTOBUF_VERSION}`).""",
87+
)
88+
89+
# Run pre-commit to cleanup the stubs
90+
subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
91+
92+
93+
if __name__ == "__main__":
94+
main()

‎scripts/sync_protobuf/s2clientprotocol.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
"""
2+
Generates the protobuf stubs for the given s2clientprotocol version using mypy-protobuf.
3+
Generally, new minor versions are a good time to update the stubs.
4+
"""
5+
6+
from __future__ import annotations
7+
8+
import re
9+
import shutil
10+
import subprocess
11+
import sys
12+
import tempfile
13+
from pathlib import Path
14+
15+
from _utils import MYPY_PROTOBUF_VERSION, REPO_ROOT, download_file, extract_archive, run_protoc, update_metadata
16+
17+
# Whenever you update PACKAGE_VERSION here, version should be updated
18+
# in stubs/s2clientprotocol/METADATA.toml and vice-versa.
19+
PACKAGE_VERSION = "c04df4adbe274858a4eb8417175ee32ad02fd609"
20+
21+
STUBS_FOLDER = REPO_ROOT / "stubs" / "s2clientprotocol"
22+
ARCHIVE_FILENAME = f"{PACKAGE_VERSION}.zip"
23+
ARCHIVE_URL = f"https://github.com/Blizzard/s2client-proto/archive/{ARCHIVE_FILENAME}"
24+
EXTRACTED_PACKAGE_DIR = f"s2client-proto-{PACKAGE_VERSION}"
25+
26+
VERSION_PATTERN = re.compile(r'def game_version\(\):\n return "(.+?)"')
27+
28+
29+
def extract_python_version(file_path: Path) -> str:
30+
"""Extract Python version from s2clientprotocol's build file"""
31+
match = re.search(VERSION_PATTERN, file_path.read_text())
32+
assert match
33+
return match.group(1)
34+
35+
36+
def main() -> None:
37+
temp_dir = Path(tempfile.mkdtemp())
38+
# Fetch s2clientprotocol (which contains all the .proto files)
39+
archive_path = temp_dir / ARCHIVE_FILENAME
40+
download_file(ARCHIVE_URL, archive_path)
41+
extract_archive(archive_path, temp_dir)
42+
43+
# Remove existing pyi
44+
for old_stub in STUBS_FOLDER.rglob("*_pb2.pyi"):
45+
old_stub.unlink()
46+
47+
PROTOC_VERSION = run_protoc(
48+
proto_paths=(EXTRACTED_PACKAGE_DIR,),
49+
mypy_out=STUBS_FOLDER,
50+
proto_globs=(f"{EXTRACTED_PACKAGE_DIR}/s2clientprotocol/*.proto",),
51+
cwd=temp_dir,
52+
)
53+
54+
PYTHON_S2_CLIENT_PROTO_VERSION = extract_python_version(temp_dir / EXTRACTED_PACKAGE_DIR / "s2clientprotocol" / "build.py")
55+
56+
# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
57+
shutil.rmtree(temp_dir)
58+
59+
update_metadata(
60+
STUBS_FOLDER,
61+
f"""Partially generated using \
62+
[mypy-protobuf=={MYPY_PROTOBUF_VERSION}](https://github.com/nipunn1313/mypy-protobuf/tree/v{MYPY_PROTOBUF_VERSION}) \
63+
and {PROTOC_VERSION} on \
64+
[s2client-proto {PYTHON_S2_CLIENT_PROTO_VERSION}](https://github.com/Blizzard/s2client-proto/tree/{PACKAGE_VERSION}).""",
65+
)
66+
67+
# Run pre-commit to cleanup the stubs
68+
subprocess.run((sys.executable, "-m", "pre_commit", "run", "--files", *STUBS_FOLDER.rglob("*_pb2.pyi")))
69+
70+
71+
if __name__ == "__main__":
72+
main()

0 commit comments

Comments
 (0)