mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-25 06:42:52 +08:00
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Generate Python protobuf stubs and make their imports package-relative."""
|
|
|
|
import argparse
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
PROTO_FILES = [
|
|
"proto/openshell.proto",
|
|
"proto/datamodel.proto",
|
|
"proto/options.proto",
|
|
"proto/sandbox.proto",
|
|
"proto/extension.proto",
|
|
"proto/credential_driver.proto",
|
|
]
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--no-pyi",
|
|
action="store_true",
|
|
help=(
|
|
"Skip .pyi type stubs. Needed with protoc older than 3.20 "
|
|
"(e.g. Fedora's grpcio-tools 1.48), which has no --pyi_out."
|
|
),
|
|
)
|
|
args = parser.parse_args()
|
|
output_dir = Path("python/openshell/_proto")
|
|
|
|
subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"grpc_tools.protoc",
|
|
"-Iproto",
|
|
"--python_out=python/openshell/_proto",
|
|
*([] if args.no_pyi else ["--pyi_out=python/openshell/_proto"]),
|
|
"--grpc_python_out=python/openshell/_proto",
|
|
*PROTO_FILES,
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
# Rewrite only imports of our generated modules, leaving google.protobuf
|
|
# imports intact. The same rule covers runtime modules and type stubs.
|
|
modules = "|".join(re.escape(f"{Path(path).stem}_pb2") for path in PROTO_FILES)
|
|
pattern = rf"^import ({modules})( as \w+)?$"
|
|
suffixes = ["_pb2.py", "_pb2_grpc.py"]
|
|
if not args.no_pyi:
|
|
suffixes.append("_pb2.pyi")
|
|
for proto in PROTO_FILES:
|
|
for suffix in suffixes:
|
|
file_path = output_dir / f"{Path(proto).stem}{suffix}"
|
|
text = file_path.read_text()
|
|
text = re.sub(pattern, r"from . import \1\2", text, flags=re.MULTILINE)
|
|
file_path.write_text(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|