| #!/usr/bin/env python | |
| import importlib.metadata | |
| import os | |
| import pathlib | |
| import grpc_tools # type: ignore | |
| from grpc_tools import protoc # type: ignore | |
| from packaging import version | |
| def get_pip_package_version(package_name: str) -> str: | |
| try: | |
| return importlib.metadata.version(package_name) | |
| except importlib.metadata.PackageNotFoundError: | |
| raise ValueError(f"Package `{package_name}` not found") | |
| protobuf_version = version.Version(get_pip_package_version("protobuf")) | |
| proto_root = os.path.join(os.path.dirname(grpc_tools.__file__), "_proto") | |
| tmp_out: pathlib.Path = pathlib.Path(f"wandb/proto/v{protobuf_version.major}/") | |
| os.chdir("../..") | |
| for proto_file in [ | |
| "wandb_base.proto", | |
| "wandb_internal.proto", | |
| "wandb_settings.proto", | |
| "wandb_telemetry.proto", | |
| "wandb_server.proto", | |
| ]: | |
| ret = protoc.main( | |
| ( | |
| "", | |
| "-I", | |
| proto_root, | |
| "-I", | |
| ".", | |
| f"--python_out={tmp_out}", | |
| f"--mypy_out={tmp_out}", | |
| f"wandb/proto/{proto_file}", | |
| ) | |
| ) | |
| assert not ret | |
| # clean up tmp dirs | |
| for p in (tmp_out / "wandb" / "proto").glob("*pb2*"): | |
| p.rename(tmp_out / p.name) | |
| os.rmdir(tmp_out / "wandb" / "proto") | |
| os.rmdir(tmp_out / "wandb") | |