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