• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
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 os
18import sys
19
20from grpc_tools import _protoc_compiler
21
22if sys.version_info >= (3, 9, 0):
23    from importlib import resources
24else:
25    import pkg_resources
26
27_PROTO_MODULE_SUFFIX = "_pb2"
28_SERVICE_MODULE_SUFFIX = "_pb2_grpc"
29
30_DISABLE_DYNAMIC_STUBS = "GRPC_PYTHON_DISABLE_DYNAMIC_STUBS"
31
32
33def main(command_arguments):
34    """Run the protocol buffer compiler with the given command-line arguments.
35
36    Args:
37      command_arguments: a list of strings representing command line arguments to
38          `protoc`.
39    """
40    command_arguments = [argument.encode() for argument in command_arguments]
41    return _protoc_compiler.run_main(command_arguments)
42
43
44def _get_resource_file_name(
45    package_or_requirement: str, resource_name: str
46) -> str:
47    """Obtain the filename for a resource on the file system."""
48    file_name = None
49    if sys.version_info >= (3, 9, 0):
50        file_name = (
51            resources.files(package_or_requirement) / resource_name
52        ).resolve()
53    else:
54        file_name = pkg_resources.resource_filename(
55            package_or_requirement, resource_name
56        )
57    return str(file_name)
58
59
60# NOTE(rbellevi): importlib.abc is not supported on 3.4.
61if sys.version_info >= (3, 5, 0):
62    import contextlib
63    import importlib
64    import importlib.abc
65    import importlib.machinery
66    import threading
67
68    _FINDERS_INSTALLED = False
69    _FINDERS_INSTALLED_LOCK = threading.Lock()
70
71    def _maybe_install_proto_finders():
72        global _FINDERS_INSTALLED
73        with _FINDERS_INSTALLED_LOCK:
74            if not _FINDERS_INSTALLED:
75                sys.meta_path.extend(
76                    [
77                        ProtoFinder(
78                            _PROTO_MODULE_SUFFIX, _protoc_compiler.get_protos
79                        ),
80                        ProtoFinder(
81                            _SERVICE_MODULE_SUFFIX,
82                            _protoc_compiler.get_services,
83                        ),
84                    ]
85                )
86
87                proto_include = _get_resource_file_name("grpc_tools", "_proto")
88                sys.path.append(proto_include)
89
90                _FINDERS_INSTALLED = True
91
92    def _module_name_to_proto_file(suffix, module_name):
93        components = module_name.split(".")
94        proto_name = components[-1][: -1 * len(suffix)]
95        # NOTE(rbellevi): The Protobuf library expects this path to use
96        # forward slashes on every platform.
97        return "/".join(components[:-1] + [proto_name + ".proto"])
98
99    def _proto_file_to_module_name(suffix, proto_file):
100        components = proto_file.split(os.path.sep)
101        proto_base_name = os.path.splitext(components[-1])[0]
102        return ".".join(components[:-1] + [proto_base_name + suffix])
103
104    def _protos(protobuf_path):
105        """Returns a gRPC module generated from the indicated proto file."""
106        _maybe_install_proto_finders()
107        module_name = _proto_file_to_module_name(
108            _PROTO_MODULE_SUFFIX, protobuf_path
109        )
110        module = importlib.import_module(module_name)
111        return module
112
113    def _services(protobuf_path):
114        """Returns a module generated from the indicated proto file."""
115        _maybe_install_proto_finders()
116        _protos(protobuf_path)
117        module_name = _proto_file_to_module_name(
118            _SERVICE_MODULE_SUFFIX, protobuf_path
119        )
120        module = importlib.import_module(module_name)
121        return module
122
123    def _protos_and_services(protobuf_path):
124        """Returns two modules, corresponding to _pb2.py and _pb2_grpc.py files."""
125        return (_protos(protobuf_path), _services(protobuf_path))
126
127    _proto_code_cache = {}
128    _proto_code_cache_lock = threading.RLock()
129
130    class ProtoLoader(importlib.abc.Loader):
131        def __init__(
132            self, suffix, codegen_fn, module_name, protobuf_path, proto_root
133        ):
134            self._suffix = suffix
135            self._codegen_fn = codegen_fn
136            self._module_name = module_name
137            self._protobuf_path = protobuf_path
138            self._proto_root = proto_root
139
140        def create_module(self, spec):
141            return None
142
143        def _generated_file_to_module_name(self, filepath):
144            components = filepath.split(os.path.sep)
145            return ".".join(
146                components[:-1] + [os.path.splitext(components[-1])[0]]
147            )
148
149        def exec_module(self, module):
150            assert module.__name__ == self._module_name
151            code = None
152            with _proto_code_cache_lock:
153                if self._module_name in _proto_code_cache:
154                    code = _proto_code_cache[self._module_name]
155                    exec(code, module.__dict__)
156                else:
157                    files = self._codegen_fn(
158                        self._protobuf_path.encode("ascii"),
159                        [path.encode("ascii") for path in sys.path],
160                    )
161                    # NOTE: The files are returned in topological order of dependencies. Each
162                    # entry is guaranteed to depend only on the modules preceding it in the
163                    # list and the last entry is guaranteed to be our requested module. We
164                    # cache the code from the first invocation at module-scope so that we
165                    # don't have to regenerate code that has already been generated by protoc.
166                    for f in files[:-1]:
167                        module_name = self._generated_file_to_module_name(
168                            f[0].decode("ascii")
169                        )
170                        if module_name not in sys.modules:
171                            if module_name not in _proto_code_cache:
172                                _proto_code_cache[module_name] = f[1]
173                            importlib.import_module(module_name)
174                    exec(files[-1][1], module.__dict__)
175
176    class ProtoFinder(importlib.abc.MetaPathFinder):
177        def __init__(self, suffix, codegen_fn):
178            self._suffix = suffix
179            self._codegen_fn = codegen_fn
180
181        def find_spec(self, fullname, path, target=None):
182            if not fullname.endswith(self._suffix):
183                return None
184            filepath = _module_name_to_proto_file(self._suffix, fullname)
185            for search_path in sys.path:
186                try:
187                    prospective_path = os.path.join(search_path, filepath)
188                    os.stat(prospective_path)
189                except (FileNotFoundError, NotADirectoryError, OSError):
190                    continue
191                else:
192                    return importlib.machinery.ModuleSpec(
193                        fullname,
194                        ProtoLoader(
195                            self._suffix,
196                            self._codegen_fn,
197                            fullname,
198                            filepath,
199                            search_path,
200                        ),
201                    )
202
203    # NOTE(rbellevi): We provide an environment variable that enables users to completely
204    # disable this behavior if it is not desired, e.g. for performance reasons.
205    if not os.getenv(_DISABLE_DYNAMIC_STUBS):
206        _maybe_install_proto_finders()
207
208if __name__ == "__main__":
209    proto_include = _get_resource_file_name("grpc_tools", "_proto")
210    sys.exit(main(sys.argv + ["-I{}".format(proto_include)]))
211