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.abc 46 import importlib.machinery 47 import threading 48 49 _FINDERS_INSTALLED = False 50 _FINDERS_INSTALLED_LOCK = threading.Lock() 51 52 def _maybe_install_proto_finders(): 53 global _FINDERS_INSTALLED 54 with _FINDERS_INSTALLED_LOCK: 55 if not _FINDERS_INSTALLED: 56 sys.meta_path.extend([ 57 ProtoFinder(_PROTO_MODULE_SUFFIX, 58 _protoc_compiler.get_protos), 59 ProtoFinder(_SERVICE_MODULE_SUFFIX, 60 _protoc_compiler.get_services) 61 ]) 62 sys.path.append( 63 pkg_resources.resource_filename('grpc_tools', '_proto')) 64 _FINDERS_INSTALLED = True 65 66 def _module_name_to_proto_file(suffix, module_name): 67 components = module_name.split(".") 68 proto_name = components[-1][:-1 * len(suffix)] 69 # NOTE(rbellevi): The Protobuf library expects this path to use 70 # forward slashes on every platform. 71 return "/".join(components[:-1] + [proto_name + ".proto"]) 72 73 def _proto_file_to_module_name(suffix, proto_file): 74 components = proto_file.split(os.path.sep) 75 proto_base_name = os.path.splitext(components[-1])[0] 76 return ".".join(components[:-1] + [proto_base_name + suffix]) 77 78 def _protos(protobuf_path): 79 """Returns a gRPC module generated from the indicated proto file.""" 80 _maybe_install_proto_finders() 81 module_name = _proto_file_to_module_name(_PROTO_MODULE_SUFFIX, 82 protobuf_path) 83 module = importlib.import_module(module_name) 84 return module 85 86 def _services(protobuf_path): 87 """Returns a module generated from the indicated proto file.""" 88 _maybe_install_proto_finders() 89 _protos(protobuf_path) 90 module_name = _proto_file_to_module_name(_SERVICE_MODULE_SUFFIX, 91 protobuf_path) 92 module = importlib.import_module(module_name) 93 return module 94 95 def _protos_and_services(protobuf_path): 96 """Returns two modules, corresponding to _pb2.py and _pb2_grpc.py files.""" 97 return (_protos(protobuf_path), _services(protobuf_path)) 98 99 _proto_code_cache = {} 100 _proto_code_cache_lock = threading.RLock() 101 102 class ProtoLoader(importlib.abc.Loader): 103 104 def __init__(self, suffix, codegen_fn, module_name, protobuf_path, 105 proto_root): 106 self._suffix = suffix 107 self._codegen_fn = codegen_fn 108 self._module_name = module_name 109 self._protobuf_path = protobuf_path 110 self._proto_root = proto_root 111 112 def create_module(self, spec): 113 return None 114 115 def _generated_file_to_module_name(self, filepath): 116 components = filepath.split(os.path.sep) 117 return ".".join(components[:-1] + 118 [os.path.splitext(components[-1])[0]]) 119 120 def exec_module(self, module): 121 assert module.__name__ == self._module_name 122 code = None 123 with _proto_code_cache_lock: 124 if self._module_name in _proto_code_cache: 125 code = _proto_code_cache[self._module_name] 126 exec(code, module.__dict__) 127 else: 128 files = self._codegen_fn( 129 self._protobuf_path.encode('ascii'), 130 [path.encode('ascii') for path in sys.path]) 131 # NOTE: The files are returned in topological order of dependencies. Each 132 # entry is guaranteed to depend only on the modules preceding it in the 133 # list and the last entry is guaranteed to be our requested module. We 134 # cache the code from the first invocation at module-scope so that we 135 # don't have to regenerate code that has already been generated by protoc. 136 for f in files[:-1]: 137 module_name = self._generated_file_to_module_name( 138 f[0].decode('ascii')) 139 if module_name not in sys.modules: 140 if module_name not in _proto_code_cache: 141 _proto_code_cache[module_name] = f[1] 142 importlib.import_module(module_name) 143 exec(files[-1][1], module.__dict__) 144 145 class ProtoFinder(importlib.abc.MetaPathFinder): 146 147 def __init__(self, suffix, codegen_fn): 148 self._suffix = suffix 149 self._codegen_fn = codegen_fn 150 151 def find_spec(self, fullname, path, target=None): 152 if not fullname.endswith(self._suffix): 153 return None 154 filepath = _module_name_to_proto_file(self._suffix, fullname) 155 for search_path in sys.path: 156 try: 157 prospective_path = os.path.join(search_path, filepath) 158 os.stat(prospective_path) 159 except (FileNotFoundError, NotADirectoryError, OSError): 160 continue 161 else: 162 return importlib.machinery.ModuleSpec( 163 fullname, 164 ProtoLoader(self._suffix, self._codegen_fn, fullname, 165 filepath, search_path)) 166 167 # NOTE(rbellevi): We provide an environment variable that enables users to completely 168 # disable this behavior if it is not desired, e.g. for performance reasons. 169 if not os.getenv(_DISABLE_DYNAMIC_STUBS): 170 _maybe_install_proto_finders() 171 172if __name__ == '__main__': 173 proto_include = pkg_resources.resource_filename('grpc_tools', '_proto') 174 sys.exit(main(sys.argv + ['-I{}'.format(proto_include)])) 175