1#! /usr/bin/env python3
2# Copyright 2021 The gRPC Authors
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15"""Builds the content of xds-protos package"""
16
17import os
18
19from grpc_tools import protoc
20import pkg_resources
21
22# We might not want to compile all the protos
23EXCLUDE_PROTO_PACKAGES_LIST = [
24    # Requires extra dependency to Prometheus protos
25    'envoy/service/metrics/v2',
26    'envoy/service/metrics/v3',
27    'envoy/service/metrics/v4alpha',
28]
29
30# Compute the pathes
31WORK_DIR = os.path.dirname(os.path.abspath(__file__))
32GRPC_ROOT = os.path.abspath(os.path.join(WORK_DIR, '..', '..', '..', '..'))
33XDS_PROTO_ROOT = os.path.join(GRPC_ROOT, 'third_party', 'envoy-api')
34UDPA_PROTO_ROOT = os.path.join(GRPC_ROOT, 'third_party', 'udpa')
35GOOGLEAPIS_ROOT = os.path.join(GRPC_ROOT, 'third_party', 'googleapis')
36VALIDATE_ROOT = os.path.join(GRPC_ROOT, 'third_party', 'protoc-gen-validate')
37OPENCENSUS_PROTO_ROOT = os.path.join(GRPC_ROOT, 'third_party',
38                                     'opencensus-proto', 'src')
39OPENTELEMETRY_PROTO_ROOT = os.path.join(GRPC_ROOT, 'third_party',
40                                        'opentelemetry')
41WELL_KNOWN_PROTOS_INCLUDE = pkg_resources.resource_filename(
42    'grpc_tools', '_proto')
43OUTPUT_PATH = WORK_DIR
44
45# Prepare the test file generation
46TEST_FILE_NAME = 'generated_file_import_test.py'
47TEST_IMPORTS = []
48
49# The pkgutil-style namespace packaging __init__.py
50PKGUTIL_STYLE_INIT = "__path__ = __import__('pkgutil').extend_path(__path__, __name__)\n"
51NAMESPACE_PACKAGES = ["google"]
52
53
54def add_test_import(proto_package_path: str,
55                    file_name: str,
56                    service: bool = False):
57    TEST_IMPORTS.append("from %s import %s\n" % (proto_package_path.replace(
58        '/', '.'), file_name.replace('.proto', '_pb2')))
59    if service:
60        TEST_IMPORTS.append("from %s import %s\n" % (proto_package_path.replace(
61            '/', '.'), file_name.replace('.proto', '_pb2_grpc')))
62
63
64# Prepare Protoc command
65COMPILE_PROTO_ONLY = [
66    'grpc_tools.protoc',
67    '--proto_path={}'.format(XDS_PROTO_ROOT),
68    '--proto_path={}'.format(UDPA_PROTO_ROOT),
69    '--proto_path={}'.format(GOOGLEAPIS_ROOT),
70    '--proto_path={}'.format(VALIDATE_ROOT),
71    '--proto_path={}'.format(WELL_KNOWN_PROTOS_INCLUDE),
72    '--proto_path={}'.format(OPENCENSUS_PROTO_ROOT),
73    '--proto_path={}'.format(OPENTELEMETRY_PROTO_ROOT),
74    '--python_out={}'.format(OUTPUT_PATH),
75]
76COMPILE_BOTH = COMPILE_PROTO_ONLY + ['--grpc_python_out={}'.format(OUTPUT_PATH)]
77
78
79def has_grpc_service(proto_package_path: str) -> bool:
80    return proto_package_path.startswith('envoy/service')
81
82
83def compile_protos(proto_root: str, sub_dir: str = '.') -> None:
84    for root, _, files in os.walk(os.path.join(proto_root, sub_dir)):
85        proto_package_path = os.path.relpath(root, proto_root)
86        if proto_package_path in EXCLUDE_PROTO_PACKAGES_LIST:
87            print(f'Skipping package {proto_package_path}')
88            continue
89        for file_name in files:
90            if file_name.endswith('.proto'):
91                # Compile proto
92                if has_grpc_service(proto_package_path):
93                    return_code = protoc.main(COMPILE_BOTH +
94                                              [os.path.join(root, file_name)])
95                    add_test_import(proto_package_path, file_name, service=True)
96                else:
97                    return_code = protoc.main(COMPILE_PROTO_ONLY +
98                                              [os.path.join(root, file_name)])
99                    add_test_import(proto_package_path,
100                                    file_name,
101                                    service=False)
102                if return_code != 0:
103                    raise Exception('error: {} failed'.format(COMPILE_BOTH))
104
105
106def create_init_file(path: str, package_path: str = "") -> None:
107    with open(os.path.join(path, "__init__.py"), 'w') as f:
108        # Apply the pkgutil-style namespace packaging, which is compatible for 2
109        # and 3. Here is the full table of namespace compatibility:
110        # https://github.com/pypa/sample-namespace-packages/blob/master/table.md
111        if package_path in NAMESPACE_PACKAGES:
112            f.write(PKGUTIL_STYLE_INIT)
113
114
115def main():
116    # Compile xDS protos
117    compile_protos(XDS_PROTO_ROOT)
118    compile_protos(UDPA_PROTO_ROOT)
119    # We don't want to compile the entire GCP surface API, just the essential ones
120    compile_protos(GOOGLEAPIS_ROOT, os.path.join('google', 'api'))
121    compile_protos(GOOGLEAPIS_ROOT, os.path.join('google', 'rpc'))
122    compile_protos(GOOGLEAPIS_ROOT, os.path.join('google', 'longrunning'))
123    compile_protos(GOOGLEAPIS_ROOT, os.path.join('google', 'logging'))
124    compile_protos(GOOGLEAPIS_ROOT, os.path.join('google', 'type'))
125    compile_protos(VALIDATE_ROOT, 'validate')
126    compile_protos(OPENCENSUS_PROTO_ROOT)
127    compile_protos(OPENTELEMETRY_PROTO_ROOT)
128
129    # Generate __init__.py files for all modules
130    create_init_file(WORK_DIR)
131    for proto_root_module in [
132            'envoy', 'google', 'opencensus', 'udpa', 'validate', 'xds',
133            'opentelemetry'
134    ]:
135        for root, _, _ in os.walk(os.path.join(WORK_DIR, proto_root_module)):
136            package_path = os.path.relpath(root, WORK_DIR)
137            create_init_file(root, package_path)
138
139    # Generate test file
140    with open(os.path.join(WORK_DIR, TEST_FILE_NAME), 'w') as f:
141        f.writelines(TEST_IMPORTS)
142
143
144if __name__ == "__main__":
145    main()
146