From 02a974c2fd8fc041f421fdc66884064bab7bae2e Mon Sep 17 00:00:00 2001 From: Dan Davison Date: Mon, 29 Jan 2024 08:51:01 -0500 Subject: [PATCH] Modularize gen-protos script (#465) --- scripts/gen_protos.py | 121 +++++++++++++++++++++++------------------- 1 file changed, 65 insertions(+), 56 deletions(-) diff --git a/scripts/gen_protos.py b/scripts/gen_protos.py index a7d50a2c..ff4ff30e 100644 --- a/scripts/gen_protos.py +++ b/scripts/gen_protos.py @@ -118,12 +118,16 @@ def fix_generated_output(base_path: Path): f.write("except ImportError:\n pass") -if __name__ == "__main__": - # Due to issues with the Python protobuf 3.x vs protobuf 4.x libraries, we - # must require that grpcio tools be on 1.48.x and protobuf be on 3.x for - # generation of protos. We can't check __version__ on the module (not - # present), and we can't use importlib.metadata due to its absence in 3.7, - # so we just run pip and check there. +def check_proto_toolchain_versions(): + """ + Check protobuf and grpcio versions. + + Due to issues with the Python protobuf 3.x vs protobuf 4.x libraries, we + must require that grpcio tools be on 1.48.x and protobuf be on 3.x for + generation of protos. We can't check __version__ on the module (not + present), and we can't use importlib.metadata due to its absence in 3.7, + so we just run pip and check there. + """ proc = subprocess.run( ["pip", "list", "--format", "freeze"], check=True, @@ -144,56 +148,61 @@ def fix_generated_output(base_path: Path): "1.48." ), f"expected 1.48.x grpcio-tools, found {grpcio_tools_version}" - print("Generating protos...", file=sys.stderr) - with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir_raw: - temp_dir = Path(temp_dir_raw) - subprocess.check_call( - [ - sys.executable, - "-mgrpc_tools.protoc", - f"--proto_path={api_proto_dir}", - f"--proto_path={core_proto_dir}", - f"--proto_path={testsrv_proto_dir}", - f"--proto_path={health_proto_dir}", - f"--proto_path={test_proto_dir}", - f"--proto_path={additional_proto_dir}", - f"--python_out={temp_dir}", - f"--grpc_python_out={temp_dir}", - f"--mypy_out={temp_dir}", - f"--mypy_grpc_out={temp_dir}", - *map(str, proto_paths), - ] + +def generate_protos(output_dir: Path): + subprocess.check_call( + [ + sys.executable, + "-mgrpc_tools.protoc", + f"--proto_path={api_proto_dir}", + f"--proto_path={core_proto_dir}", + f"--proto_path={testsrv_proto_dir}", + f"--proto_path={health_proto_dir}", + f"--proto_path={test_proto_dir}", + f"--proto_path={additional_proto_dir}", + f"--python_out={output_dir}", + f"--grpc_python_out={output_dir}", + f"--mypy_out={output_dir}", + f"--mypy_grpc_out={output_dir}", + *map(str, proto_paths), + ] + ) + # Remove every _grpc.py file that isn't part of a Temporal "service" + for grpc_file in output_dir.glob("**/*_grpc.py*"): + if ( + len(grpc_file.parents) < 2 + or grpc_file.parents[0].name != "v1" + or not grpc_file.parents[1].name.endswith("service") + ): + grpc_file.unlink() + # Apply fixes before moving code + fix_generated_output(output_dir) + # Move protos + for p in (output_dir / "temporal" / "api").iterdir(): + shutil.rmtree(api_out_dir / p.name, ignore_errors=True) + p.replace(api_out_dir / p.name) + shutil.rmtree(api_out_dir / "dependencies", ignore_errors=True) + for p in (output_dir / "temporal" / "sdk" / "core").iterdir(): + shutil.rmtree(sdk_out_dir / p.name, ignore_errors=True) + p.replace(sdk_out_dir / p.name) + shutil.rmtree(sdk_out_dir / "health", ignore_errors=True) + (output_dir / "health").replace(sdk_out_dir / "health") + # Move test protos + for v in ["__init__.py", "proto_message_pb2.py", "proto_message_pb2.pyi"]: + shutil.copy2( + output_dir / "worker" / "workflow_sandbox" / "testmodules" / "proto" / v, + test_proto_dir + / "worker" + / "workflow_sandbox" + / "testmodules" + / "proto" + / v, ) - # Remove every _grpc.py file that isn't part of a Temporal "service" - for grpc_file in temp_dir.glob("**/*_grpc.py*"): - if ( - len(grpc_file.parents) < 2 - or grpc_file.parents[0].name != "v1" - or not grpc_file.parents[1].name.endswith("service") - ): - grpc_file.unlink() - # Apply fixes before moving code - fix_generated_output(temp_dir) - # Move protos - for p in (temp_dir / "temporal" / "api").iterdir(): - shutil.rmtree(api_out_dir / p.name, ignore_errors=True) - p.replace(api_out_dir / p.name) - shutil.rmtree(api_out_dir / "dependencies", ignore_errors=True) - for p in (temp_dir / "temporal" / "sdk" / "core").iterdir(): - shutil.rmtree(sdk_out_dir / p.name, ignore_errors=True) - p.replace(sdk_out_dir / p.name) - shutil.rmtree(sdk_out_dir / "health", ignore_errors=True) - (temp_dir / "health").replace(sdk_out_dir / "health") - # Move test protos - for v in ["__init__.py", "proto_message_pb2.py", "proto_message_pb2.pyi"]: - shutil.copy2( - temp_dir / "worker" / "workflow_sandbox" / "testmodules" / "proto" / v, - test_proto_dir - / "worker" - / "workflow_sandbox" - / "testmodules" - / "proto" - / v, - ) + +if __name__ == "__main__": + check_proto_toolchain_versions() + print("Generating protos...", file=sys.stderr) + with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir: + generate_protos(Path(temp_dir)) print("Done", file=sys.stderr)