protogen

#!/usr/bin/env python3
#
# Copyright 2018 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------------

from glob import glob
import os
import re
import tempfile
from os.path import dirname as DIRNAME
import sys

from grpc.tools.protoc import main as _protoc


JOIN = os.path.join
TOP_DIR = os.path.dirname(os.path.realpath(__file__))


def protoc_go(src_dir, base_dir, pkg):
    # 1. Get a list of all .proto files to build
    cwd = os.getcwd()
    os.chdir(src_dir)
    proto_files = glob("*.proto")
    os.chdir(cwd)
    # 2. Create a temp directory for building
    with tempfile.TemporaryDirectory() as tmp_dir:
        defer = []
        # 3. In order to have each protobuf file in its own namespace,
        # directories need to be created for all of them and the .proto file
        # copied in
        for proto_file in proto_files:
            proto = proto_file[:-6]
            sub_pkg = JOIN(pkg, proto)
            tmp_pkg_dir = JOIN(tmp_dir, sub_pkg + "_pb2")
            os.makedirs(tmp_pkg_dir)
            src = JOIN(src_dir, proto_file)
            dst = JOIN(tmp_pkg_dir, proto_file)
            with open(src, encoding='utf-8') as fin:
                with open(dst, "w", encoding='utf-8') as fout:
                    src_contents = fin.read()
                    fixed_contents = fix_import(src_contents, pkg, sub_dir=True)
                    fout.write(fixed_contents)
            # Need to defer until the whole directory is setup
            defer.append([
                __file__,
                "-I=%s" % tmp_dir,
                "--go_out=%s" % JOIN(TOP_DIR, base_dir),
                "%s" % dst
            ])
        for args in defer:
            _protoc(args)


def fix_import(contents, pkg, sub_dir=False):
    pattern = r'^import "(.*)\.proto\"'
    if sub_dir:
        template = r'import "%s/\1_pb2/\1.proto"'
    else:
        template = r'import "%s/\1.proto"'
    return re.sub(
        pattern,
        lambda match: match.expand(template) % pkg,
        contents,
        flags=re.MULTILINE)


def main():
    protoc_go("protos", "./", "protobuf")
    protoc_go("examples/smallbank/protos",
              "examples/smallbank/smallbank_go/src",
              "protobuf")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        exit(1)