Rewrote protobuf generation scripts in Python #12527

Avasam · 2024-08-15T06:13:44Z

This is so much faster on Windows, takes a few seconds to run (including downloads and pre-commit). Compared to the over 1m it used to take me on WSL.

I have two design questions:

Which usage do we prefer:

# A)
python scripts/sync_proto/tensorflow.py

# B)
python scripts/sync_proto.py tensorflow
# Which I could make work with the following for shell autocomplete
python scripts/sync_proto.py ./stubs/tensorflow

Should I add a parametrized main function (something like sync_stubs_with_proto) that takes all of a scripts' special needs as parameters (including a "post-run" Callable). Or keep the 3 scripts separate with shared helper functions.

I'm also open to name changes suggestions.

Avasam · 2024-08-16T19:20:57Z

requirements-tests.txt

@@ -13,6 +13,8 @@ ruff==0.5.4              # must match .pre-commit-config.yaml

 # Libraries used by our various scripts.
 aiohttp==3.10.2
+grpcio-tools


I'm sure there's a minimal version of protoc that is shipped with grpcio-tools that we should be using, but I can't recall off the top of my head and would have to search through past PRs to find what it was.

Avasam · 2024-08-16T19:21:48Z

scripts/sync_protobuf/_helpers.py

+from http.client import HTTPResponse
+from pathlib import Path
+from typing import TYPE_CHECKING, Iterable
+from urllib.request import urlopen


I'm purposefully avoiding requests here, as to not add requests and types-requests in requirements-tests.txt

Avasam · 2024-08-16T19:24:12Z

scripts/sync_protobuf/google_protobuf.py

+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)


Technically the type is dict[str, dict[str, str | dict[str, str]]] so I'm being a bit cheesy here, but I think it's fine over a Any. This is the type for the tree of key-values we traverse, everything else doesn't matter. A TypedDict felt overkill just to prevent a typo in data.values().

…to proto-scripts

Avasam · 2024-08-16T19:36:16Z

requirements-tests.txt

+# 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"


xref grpc/grpc#36201 & grpc/grpc#34922

…scripts

srittau

Thanks! I'm not really familiar with protobuf or protobuf stubs generation so my comments are limited to general issues.

In general, there are a few places that use / as path separator, so I'd expect problems on Windows, but I'm fine with that for now.

srittau · 2024-09-12T09:09:16Z

scripts/sync_protobuf/_utils.py

+    resp: HTTPResponse = urlopen(url)
+    if resp.getcode() != 200:
+        raise RuntimeError(f"Error downloading {url}")
+    with open(destination, "wb") as file:
+        file.write(resp.read())


Suggested change

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())

srittau · 2024-09-12T09:17:59Z

scripts/sync_protobuf/google_protobuf.py

+        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"]


I'd like to see some validation of the version, considering its coming from an outside source. Something like:

Suggested change

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"]

data = json.load(file)

# The root key will be the protobuf source code version

version = next(iter(data.values()))["languages"]["python"]

assert isinstance(version, str)

assert re.fullmatch(r"...", version) # proper re here

return version

This way we're also sure (at runtime) that version has the correct type and format.

srittau · 2024-09-12T09:21:34Z

scripts/sync_protobuf/google_protobuf.py

+    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()))


Nit:

Suggested change

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

srittau · 2024-09-12T09:29:06Z

scripts/sync_protobuf/google_protobuf.py

+    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)


To make sure the temp directory is always cleaned up:

Suggested change

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)

with tempfile.TemporaryDirectory() as td:

temp_dir = Path(td)

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

srittau · 2024-09-12T09:30:03Z

scripts/sync_protobuf/s2clientprotocol.py

+
+
+def main() -> None:
+    temp_dir = Path(tempfile.mkdtemp())


srittau · 2024-09-12T09:32:16Z

scripts/sync_protobuf/tensorflow.py

+
+
+def main() -> None:
+    temp_dir = Path(tempfile.mkdtemp())


Avasam added 5 commits August 15, 2024 02:12

WIP (s2clientprotocol done)

db53bda

Completed tensorflow script

cdd58f3

Fix CI tests

9f33b0e

Finished google protobuf

c1ba6e5

Merge branch 'main' into proto-scripts

da40ac3

Avasam marked this pull request as ready for review August 16, 2024 19:17

Avasam commented Aug 16, 2024

View reviewed changes

Forgot to delete old script

5f2ec26

Avasam changed the title ~~[WIP] Rewrote protobuf generation scripts in Python~~ Rewrote protobuf generation scripts in Python Aug 16, 2024

Avasam added 2 commits August 16, 2024 15:35

Don't install grpcio-tools on Python 3.13

893a223

Merge branch 'proto-scripts' of https://github.com/Avasam/typeshed in…

b23bbae

…to proto-scripts

Avasam commented Aug 16, 2024

View reviewed changes

Avasam added the project: infrastructure typeshed build, test, documentation, or distribution related label Aug 18, 2024

Avasam added 2 commits August 19, 2024 11:47

Merge branch 'main' of https://github.com/python/typeshed into proto-…

b0a2250

…scripts

Rename _helpers to _utils

75a746c

Avasam mentioned this pull request Sep 7, 2024

[stubsabot] Bump protobuf to 5.28.* #12604

Open

Merge branch 'main' into proto-scripts

d00e261

srittau requested changes Sep 12, 2024

View reviewed changes

Avasam added the status: needs decision Needs a final decision by the typeshed maintainers label Sep 12, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Rewrote protobuf generation scripts in Python #12527

Rewrote protobuf generation scripts in Python #12527

Avasam commented Aug 15, 2024 •

edited

Loading

Avasam Aug 16, 2024

Avasam Aug 16, 2024 •

edited

Loading

Avasam Aug 16, 2024 •

edited

Loading

Avasam Aug 16, 2024

srittau left a comment

srittau Sep 12, 2024

srittau Sep 12, 2024

srittau Sep 12, 2024

srittau Sep 12, 2024

srittau Sep 12, 2024

srittau Sep 12, 2024

-        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"]
+        data = json.load(file)
+        # The root key will be the protobuf source code version
+        version = next(iter(data.values()))["languages"]["python"]
+        assert isinstance(version, str)
+        assert re.fullmatch(r"...", version)  # proper re here
+        return version

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

Rewrote protobuf generation scripts in Python #12527

Are you sure you want to change the base?

Rewrote protobuf generation scripts in Python #12527

Conversation

Avasam commented Aug 15, 2024 • edited Loading

Choose a reason for hiding this comment

Avasam Aug 16, 2024 • edited Loading

Choose a reason for hiding this comment

Avasam Aug 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

srittau left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Avasam commented Aug 15, 2024 •

edited

Loading

Avasam Aug 16, 2024 •

edited

Loading

Avasam Aug 16, 2024 •

edited

Loading