File size: 1,278 Bytes
9c6594c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#!/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")