• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2
3# Copyright 2016 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import pkg_resources
18import sys
19
20import os
21
22from grpc_tools import _protoc_compiler
23
24_PROTO_MODULE_SUFFIX = "_pb2"
25_SERVICE_MODULE_SUFFIX = "_pb2_grpc"
26
27_DISABLE_DYNAMIC_STUBS = "GRPC_PYTHON_DISABLE_DYNAMIC_STUBS"
28
29
30def main(command_arguments):
31    """Run the protocol buffer compiler with the given command-line arguments.
32
33  Args:
34    command_arguments: a list of strings representing command line arguments to
35        `protoc`.
36  """
37    command_arguments = [argument.encode() for argument in command_arguments]
38    return _protoc_compiler.run_main(command_arguments)
39
40
41# NOTE(rbellevi): importlib.abc is not supported on 3.4.
42if sys.version_info >= (3, 5, 0):
43    import contextlib
44    import importlib
45    import importlib.machinery
46    import threading
47
48    _FINDERS_INSTALLED = False
49    _FINDERS_INSTALLED_LOCK = threading.Lock()
50
51    def _maybe_install_proto_finders():
52        global _FINDERS_INSTALLED
53        with _FINDERS_INSTALLED_LOCK:
54            if not _FINDERS_INSTALLED:
55                sys.meta_path.extend([
56                    ProtoFinder(_PROTO_MODULE_SUFFIX,
57                                _protoc_compiler.get_protos),
58                    ProtoFinder(_SERVICE_MODULE_SUFFIX,
59                                _protoc_compiler.get_services)
60                ])
61                sys.path.append(
62                    pkg_resources.resource_filename('grpc_tools', '_proto'))
63                _FINDERS_INSTALLED = True
64
65    def _module_name_to_proto_file(suffix, module_name):
66        components = module_name.split(".")
67        proto_name = components[-1][:-1 * len(suffix)]
68        # NOTE(rbellevi): The Protobuf library expects this path to use
69        # forward slashes on every platform.
70        return "/".join(components[:-1] + [proto_name + ".proto"])
71
72    def _proto_file_to_module_name(suffix, proto_file):
73        components = proto_file.split(os.path.sep)
74        proto_base_name = os.path.splitext(components[-1])[0]
75        return ".".join(components[:-1] + [proto_base_name + suffix])
76
77    def _protos(protobuf_path):
78        """Returns a gRPC module generated from the indicated proto file."""
79        _maybe_install_proto_finders()
80        module_name = _proto_file_to_module_name(_PROTO_MODULE_SUFFIX,
81                                                 protobuf_path)
82        module = importlib.import_module(module_name)
83        return module
84
85    def _services(protobuf_path):
86        """Returns a module generated from the indicated proto file."""
87        _maybe_install_proto_finders()
88        _protos(protobuf_path)
89        module_name = _proto_file_to_module_name(_SERVICE_MODULE_SUFFIX,
90                                                 protobuf_path)
91        module = importlib.import_module(module_name)
92        return module
93
94    def _protos_and_services(protobuf_path):
95        """Returns two modules, corresponding to _pb2.py and _pb2_grpc.py files."""
96        return (_protos(protobuf_path), _services(protobuf_path))
97
98    _proto_code_cache = {}
99    _proto_code_cache_lock = threading.RLock()
100
101    class ProtoLoader(importlib.abc.Loader):
102
103        def __init__(self, suffix, codegen_fn, module_name, protobuf_path,
104                     proto_root):
105            self._suffix = suffix
106            self._codegen_fn = codegen_fn
107            self._module_name = module_name
108            self._protobuf_path = protobuf_path
109            self._proto_root = proto_root
110
111        def create_module(self, spec):
112            return None
113
114        def _generated_file_to_module_name(self, filepath):
115            components = filepath.split(os.path.sep)
116            return ".".join(components[:-1] +
117                            [os.path.splitext(components[-1])[0]])
118
119        def exec_module(self, module):
120            assert module.__name__ == self._module_name
121            code = None
122            with _proto_code_cache_lock:
123                if self._module_name in _proto_code_cache:
124                    code = _proto_code_cache[self._module_name]
125                    exec(code, module.__dict__)
126                else:
127                    files = self._codegen_fn(
128                        self._protobuf_path.encode('ascii'),
129                        [path.encode('ascii') for path in sys.path])
130                    # NOTE: The files are returned in topological order of dependencies. Each
131                    # entry is guaranteed to depend only on the modules preceding it in the
132                    # list and the last entry is guaranteed to be our requested module. We
133                    # cache the code from the first invocation at module-scope so that we
134                    # don't have to regenerate code that has already been generated by protoc.
135                    for f in files[:-1]:
136                        module_name = self._generated_file_to_module_name(
137                            f[0].decode('ascii'))
138                        if module_name not in sys.modules:
139                            if module_name not in _proto_code_cache:
140                                _proto_code_cache[module_name] = f[1]
141                            importlib.import_module(module_name)
142                    exec(files[-1][1], module.__dict__)
143
144    class ProtoFinder(importlib.abc.MetaPathFinder):
145
146        def __init__(self, suffix, codegen_fn):
147            self._suffix = suffix
148            self._codegen_fn = codegen_fn
149
150        def find_spec(self, fullname, path, target=None):
151            if not fullname.endswith(self._suffix):
152                return None
153            filepath = _module_name_to_proto_file(self._suffix, fullname)
154            for search_path in sys.path:
155                try:
156                    prospective_path = os.path.join(search_path, filepath)
157                    os.stat(prospective_path)
158                except (FileNotFoundError, NotADirectoryError, OSError):
159                    continue
160                else:
161                    return importlib.machinery.ModuleSpec(
162                        fullname,
163                        ProtoLoader(self._suffix, self._codegen_fn, fullname,
164                                    filepath, search_path))
165
166    # NOTE(rbellevi): We provide an environment variable that enables users to completely
167    # disable this behavior if it is not desired, e.g. for performance reasons.
168    if not os.getenv(_DISABLE_DYNAMIC_STUBS):
169        _maybe_install_proto_finders()
170
171if __name__ == '__main__':
172    proto_include = pkg_resources.resource_filename('grpc_tools', '_proto')
173    sys.exit(main(sys.argv + ['-I{}'.format(proto_include)]))
174