1# Copyright (C) 2016 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14import logging
15import os
16import shutil
17import subprocess
18import sys
19from importlib import import_module
20
21from google.protobuf import descriptor_pb2
22from google.protobuf import text_format
23
24
25def compile_proto(proto_path, output_dir):
26    """Invoke Protocol Compiler to generate python from given source .proto."""
27    # Find compiler path
28    protoc = None
29    if 'PROTOC' in os.environ and os.path.exists(os.environ['PROTOC']):
30        protoc = os.environ['PROTOC']
31    if not protoc:
32        protoc = shutil.which('protoc')
33    if not protoc:
34        logging.error(
35            "Cannot find protobuf compiler (>=3.0.0), please install"
36            "protobuf-compiler package. Prefer copying from <top>/prebuilts/tools"
37        )
38        logging.error("    prebuilts/tools/linux-x86_64/protoc/bin/protoc")
39        logging.error("If prebuilts are not available, use apt-get:")
40        logging.error("    sudo apt-get install protobuf-compiler")
41        return None
42    # Validate input proto path
43    if not os.path.exists(proto_path):
44        logging.error('Can\'t find required file: %s\n' % proto_path)
45        return None
46    # Validate output py-proto path
47    if not os.path.exists(output_dir):
48        os.makedirs(output_dir)
49    elif not os.path.isdir(output_dir):
50        logging.error("Output path is not a valid directory: %s" %
51                      (output_dir))
52        return None
53    input_dir = os.path.dirname(proto_path)
54    output_filename = os.path.basename(proto_path).replace('.proto', '_pb2.py')
55    output_path = os.path.join(output_dir, output_filename)
56    # Compiling proto
57    logging.debug('Generating %s' % output_path)
58    protoc_command = [
59        protoc, '-I=%s' % (input_dir), '--python_out=%s' % (output_dir),
60        proto_path
61    ]
62    logging.debug('Running command %s' % protoc_command)
63    if subprocess.call(protoc_command, stderr=subprocess.STDOUT) != 0:
64        logging.error("Fail to compile proto")
65        return None
66    output_module_name = os.path.splitext(output_filename)[0]
67    return output_module_name
68
69
70def compile_import_proto(output_dir, proto_path):
71    """Compiles the given protobuf file and return the module.
72
73    Args:
74        output_dir: The directory to put the compilation output.
75        proto_path: The path to the .proto file that needs to be compiled.
76    Returns:
77        The protobuf module.
78    """
79    output_module_name = compile_proto(proto_path, output_dir)
80    if not output_module_name:
81        return None
82    sys.path.append(output_dir)
83    output_module = None
84    try:
85        output_module = import_module(output_module_name)
86    except ImportError:
87        logging.error("Cannot import generated py-proto %s" %
88                      (output_module_name))
89    return output_module
90
91
92def parse_proto_to_ascii(binary_proto_msg):
93    """ Parses binary protobuf message to human readable ascii string.
94
95    Args:
96        binary_proto_msg: The binary format of the proto message.
97    Returns:
98        The ascii format of the proto message.
99    """
100    return text_format.MessageToString(binary_proto_msg)
101
102
103def to_descriptor_proto(proto):
104    """Retrieves the descriptor proto for the given protobuf message.
105
106    Args:
107        proto: the original message.
108    Returns:
109        The descriptor proto for the input meessage.
110    """
111    descriptor_proto = descriptor_pb2.DescriptorProto()
112    proto.DESCRIPTOR.CopyToProto(descriptor_proto)
113    return descriptor_proto
114