generate_infra_protos.py

#!/usr/bin/env python
import os
import shutil
import subprocess

current_file_dir = os.path.dirname(os.path.abspath(__file__))

protoBaseDir = "scheduke/luci_proto/go.chromium.org/chromiumos/"
chrootSide = "proto"
protoDir = os.path.join(current_file_dir, protoBaseDir)
chrootProtoDir = os.path.join(protoDir, chrootSide)
schedukeDir = os.path.join(current_file_dir, "scheduke/")
generatedDir = os.path.join(schedukeDir, "chromiumos")

updateFile = os.path.join(current_file_dir,
                          os.path.join(protoBaseDir, "update.py"))
# Remove old protos
for file in os.listdir(os.path.dirname(updateFile)):
  if file.endswith("update.py"):
    continue
  shutil.rmtree(os.path.join(os.path.dirname(updateFile), file))
  print(file)
subprocess.call(f"python3 {updateFile}", shell=True)

# Discover protos to protoc
dirsToBuild = set()
for root, dirs, files in os.walk(protoDir, topdown=False):
  for name in files:
    if name.endswith(".proto"):
      dirsToBuild.add(root)

# Discover protos to GRPC protoc
grpcDirsToBuild = set()
for root, dirs, files in os.walk(protoDir, topdown=False):
  for name in files:
    if name.endswith("_service.proto"):
      grpcDirsToBuild.add(root)

# # nuke old bindings.
if os.path.isdir(os.path.join(generatedDir)):
  shutil.rmtree(os.path.join(generatedDir))
os.mkdir(generatedDir)

buildStr = "python3 -m grpc_tools.protoc -I {chrootProtoDir} --proto_path={protoDir} --python_out={schedukeDir} --pyi_out={schedukeDir} {path}/*.proto"
buildSchedukeStr = "python -m grpc_tools.protoc -I {chrootProtoDir} --proto_path={protoDir} --python_out={schedukeDir} --pyi_out={schedukeDir} " \
                   "--go_out={schedukeDir}/event_puller/service {path}/*.proto"
grpcBuildStr = "python3 -m grpc_tools.protoc -I {chrootProtoDir} --proto_path={protoDir} --grpc_python_out={schedukeDir} {path}/*_service.proto"

# Protoc them.
for path in dirsToBuild:
  protoBuildStr = buildStr.format(chrootProtoDir=chrootProtoDir,
                                  protoDir=protoDir, path=path,
                                  schedukeDir=schedukeDir)
  # Generate go bindings for Scheduke-specific protos only.
  if "chromiumos/test/scheduling" in path:
    protoBuildStr = buildSchedukeStr.format(chrootProtoDir=chrootProtoDir,
                                            protoDir=protoDir, path=path,
                                            schedukeDir=schedukeDir,
                                            current_file_dir=current_file_dir)
  returned_value = subprocess.call(protoBuildStr, shell=True)
  if returned_value != 0:
    print(f"Error building {protoBuildStr}")
    exit(1)

# GRPC protoc them.
for path in grpcDirsToBuild:
  grpcProtoBuildStr = grpcBuildStr.format(chrootProtoDir=chrootProtoDir,
                                          protoDir=protoDir, path=path,
                                          schedukeDir=schedukeDir)
  returned_value = subprocess.call(grpcProtoBuildStr, shell=True)
  if returned_value != 0:
    print(f"Error building {grpcProtoBuildStr}")
    exit(1)

print('done. probably.')