import pathlib import os import sys import subprocess import configparser import dataclasses from typing import List CONFIG_PATH: pathlib.Path = pathlib.Path("./setup.cfg").absolute() @dataclasses.dataclass class Options: """Dataclass used to hold the relevant options from our config file.""" proto_root_dir: pathlib.Path """Root path to our protobuf folder.""" go_module_root: pathlib.Path """Root module to use for protoc-gen-go '--go_opt=module=' flag.""" def load_cfg() -> Options: """ loads library config file :return: loaded `Options` object """ config = configparser.ConfigParser() config.read(str(CONFIG_PATH)) options = Options( proto_root_dir=pathlib.Path(config["proto"]["root_source_path"]), go_module_root=pathlib.Path(config["proto"]["root_go_package"]), ) return options def main() -> None: """Run the script.""" options = load_cfg() generate_golang_source(options) add_bson_tags() def generate_golang_source(options: Options) -> None: """Generate the protocol buffers.""" proto_files = find_proto_files(options) run_protoc_command(proto_files, options) def find_proto_files(options: Options) -> List[str]: """Glob all the proto files in the root proto folder and return as list.""" proto_file_list: List[str] = list() for proto_path in options.proto_root_dir.rglob("./**/*.proto"): if "google" in str(proto_path): continue path_str = str(proto_path) path_str = path_str.replace(str(os.getcwd()), ".") proto_file_list.append(path_str) return proto_file_list def run_protoc_command(proto_files: List[str], options: Options) -> None: """Run the protoc command and stream the output to std out.""" command = build_protoc_command(proto_files, options) proc = subprocess.Popen(command) _, _ = proc.communicate() if proc.returncode != 0: sys.exit(proc.returncode) def build_protoc_command(protoc_files: List[str], options: Options) -> List[str]: """Put together the protoc command to buid.""" command = [ "protoc", "--experimental_allow_proto3_optional", "--go_out=plugins=grpc:.", f"--go_opt=module={options.go_module_root}", ] command.extend(protoc_files) return command def add_bson_tags() -> None: """Add bson tags through protoc-go-inject-tag.""" directory = pathlib.Path(os.getcwd()) for source_code_file_path in directory.rglob("./**/*.pb.go"): run_tag_command(source_code_file_path) def build_tag_command(source_code_file_path: pathlib.Path) -> List[str]: """Build the protoc-go-inject-tag command.""" return [ "protoc-go-inject-tag", f"-input={source_code_file_path}", "-XXX_skip=bson", ] def run_tag_command(source_code_path: pathlib.Path) -> None: """Run the protoc-go-inject-tag command and stream the output to stdout.""" command = build_tag_command(source_code_path) proc = subprocess.Popen(command) _, _ = proc.communicate() if proc.returncode != 0: sys.exit(proc.returncode) if __name__ == "__main__": main()