compile_single.py

#!/usr/bin/env python

import argparse
import os
import shutil
import re

from helpers import compile_helper
from helpers import go_helper

from subprocess import call

# Add this to your path
protoc_path = "protoc"

# Specify desired language / output
parser = argparse.ArgumentParser()
parser.add_argument("-l", "--lang", help="Language to produce protoc files")
parser.add_argument("-o", "--out_path", help="Output path for protoc files")
parser.add_argument("-d", "--desc_file", action='store_true', help="For generating a .desc file only")
parser.add_argument("--go_import_prefix", help="Prefix all imports in output go files for vendoring all dependencies")
parser.add_argument("--go_package", help="The name of the exported go package")
parser.add_argument("--java_multiple_files", action='store_true', help="Write each message to a separate .java file.")
args = parser.parse_args()

# Set defaults
lang = args.lang or "csharp"
out_path = args.out_path or "out"
desc_file = args.desc_file
default_out_path = out_path == "out"
go_import_prefix = args.go_import_prefix
go_package = args.go_package or "protos"
java_multiple_files = args.java_multiple_files

# Determine where to store
proto_path = os.path.abspath("src")
tmp_path = os.path.abspath("tmp")
out_path = os.path.abspath(out_path)

# Clean up previous
if os.path.exists(tmp_path):
    shutil.rmtree(tmp_path)

if default_out_path and os.path.exists(out_path):
    shutil.rmtree(out_path)

# Create necessary directory
os.makedirs(tmp_path)

if not os.path.exists(out_path):
    os.makedirs(out_path)

created_packages = []
package_mappings = []

# Go specific
go_package_mappings = []


def get_package(path):
    for file_name in os.listdir(path):
        file_name_path = os.path.join(path, file_name)
        if os.path.isfile(file_name_path) and file_name.endswith('.proto'):
            with open(file_name_path, 'r') as proto_file:
                for proto_line in proto_file.readlines():
                    if proto_line.startswith("package"):
                        return re.search('package (.*?);', proto_line).group(1)
    return None


def walk_files(main_file, path, package, imports=None):
    if imports is None:
        imports = []

    if not desc_file and package == "POGOProtos":
        print("Can't compile..")
        print("File: '%s'" % path)
        print("Please place the file in 'src/POGOProtos/' in a sub-directory.")
        exit()

    main_file.write('syntax = "proto3";\n')

    short_package_name = str.split(package, '.')[-1].lower()

    main_file.write('package %s;\n\n' % package)

    if lang == "go":
        package = go_helper.convert_to_go_package(package)
        main_file.write('option go_package = "%s";\n' % go_package)

    if java_multiple_files:
        main_file.write('option java_multiple_files = true;\n')

    messages = ""

    for file_name in os.listdir(path):
        file_name_path = os.path.join(path, file_name)
        if file_name_path.endswith(".proto") and os.path.isfile(file_name_path):
            with open(file_name_path, 'r') as proto_file:
                is_header = True
                for proto_line in proto_file.readlines():
                    if proto_line.startswith("message") or proto_line.startswith("enum"):
                        is_header = False

                    if is_header:
                        if proto_line.startswith("import"):
                            import_from_package_re = re.search('import (public )?"(.*?)(\/)?([a-zA-Z0-9]+\.proto)";',
                                                               proto_line)

                            if import_from_package_re is None:
                                print("Can't compile..")
                                print("File: '%s'" % file_name_path)
                                print("Bad import line: '%s'" % proto_line)
                                exit()

                            import_from_package = import_from_package_re.group(2).replace("/", ".")

                            if lang == "go":
                                import_from_package = go_helper.convert_to_go_package(import_from_package)

                            if import_from_package not in imports:
                                imports.append(import_from_package)

                    if not is_header:
                        messages += proto_line

                        if proto_line == "}":
                            messages += "\n"

    for package_import in imports:
        if package_import != package:
            main_file.write('import public "%s.proto";\n' % package_import)

    if len(imports) is not 0:
        main_file.write('\n')

    main_file.writelines(messages)


def walk_directory(path):
    for dir_name in os.listdir(path):
        dir_name_path = os.path.join(path, dir_name)
        if os.path.isdir(dir_name_path):
            package = get_package(dir_name_path)
            if package is not None:
                if lang == "go":
                    file_name = go_helper.convert_to_go_package(package)
                    package_mappings.append([file_name, (file_name + ".proto")])
                else:
                    file_name = package
                    package_mappings.append([file_name, (file_name + ".proto")])

                package_file_path = os.path.join(tmp_path, file_name + ".proto")

                if lang == "go":
                    package_directory = os.path.dirname(package_file_path)
                    compile_helper.mkdir_p(package_directory)

                with open(package_file_path, 'a') as package_file:
                    walk_files(package_file, dir_name_path, package)
                    created_packages.append(package)

            walk_directory(dir_name_path)


def compile_go_package(path):
    proto_files = compile_helper.abslistdir(path)

    # Compile with the grpc plugin
    command_out_path = "plugins=grpc"

    # Allow to specify import_prefix for complete vendoring of dependencies
    if go_import_prefix:
        command_out_path += ",import_prefix=%s" % go_import_prefix

    # Combine the output with all other output options
    command_out_path = "%s:%s" % (command_out_path, os.path.abspath(out_path))

    command = """{0} --proto_path="{1}" --go_out={2} {3}""".format(
        protoc_path,
        tmp_path,
        command_out_path,
        proto_files
    )

    call(command, shell=True)


def compile_directories(path):
    for proto_file_name in os.listdir(path):
        command_out_path = os.path.abspath(out_path)
        item_path = os.path.join(path, proto_file_name)

        if os.path.isfile(item_path):

            command = """{0} --proto_path="{1}" --{2}_out="{3}" "{4}\"""".format(
                protoc_path,
                tmp_path,
                lang,
                command_out_path,
                item_path
            )

            call(command, shell=True)

        elif os.path.isdir(item_path):
            compile_directories(item_path)


walk_directory(proto_path)

# Compile =)
if desc_file:
    root_package_file_path = os.path.join(tmp_path, "POGOProtos.proto")
    with open(root_package_file_path, 'a') as root_package_file:
        walk_files(root_package_file, proto_path, "POGOProtos", created_packages)

    command = """{0} --include_imports --proto_path="{1}" --descriptor_set_out="{2}" "{3}\"""".format(
        protoc_path,
        tmp_path,
        os.path.abspath(out_path + "/POGOProtos.desc"),
        os.path.join(tmp_path, "POGOProtos.proto")
    )

    call(command, shell=True)
else:
    if lang == "go":
        compile_go_package(tmp_path)
    else:
        compile_directories(tmp_path)

    compile_helper.finish_compile(out_path, lang)

print("Done!")