|
|
|
@ -19,9 +19,9 @@ import sys |
|
|
|
|
|
|
|
|
|
# TODO: Figure out how to add this dependency to setuptools. |
|
|
|
|
import six |
|
|
|
|
import imp |
|
|
|
|
import os |
|
|
|
|
|
|
|
|
|
import contextlib |
|
|
|
|
import importlib |
|
|
|
|
import importlib.machinery |
|
|
|
|
import sys |
|
|
|
@ -38,53 +38,45 @@ def main(command_arguments): |
|
|
|
|
command_arguments = [argument.encode() for argument in command_arguments] |
|
|
|
|
return _protoc_compiler.run_main(command_arguments) |
|
|
|
|
|
|
|
|
|
def _module_name_to_proto_file(module_name): |
|
|
|
|
def _module_name_to_proto_file(suffix, module_name): |
|
|
|
|
components = module_name.split(".") |
|
|
|
|
proto_name = components[-1][:-1*len("_pb2")] |
|
|
|
|
proto_name = components[-1][:-1*len(suffix)] |
|
|
|
|
return os.path.sep.join(components[:-1] + [proto_name + ".proto"]) |
|
|
|
|
|
|
|
|
|
def _proto_file_to_module_name(proto_file): |
|
|
|
|
def _proto_file_to_module_name(suffix, proto_file): |
|
|
|
|
components = proto_file.split(os.path.sep) |
|
|
|
|
proto_base_name = os.path.splitext(components[-1])[0] |
|
|
|
|
return os.path.sep.join(components[:-1] + [proto_base_name + "_pb2"]) |
|
|
|
|
|
|
|
|
|
def _import_modules_from_files(files): |
|
|
|
|
modules = [] |
|
|
|
|
for filename, code in files: |
|
|
|
|
base_name = os.path.basename(filename.decode('ascii')) |
|
|
|
|
proto_name, _ = os.path.splitext(base_name) |
|
|
|
|
anchor_package = ".".join(os.path.normpath(os.path.dirname(filename.decode('ascii'))).split(os.sep)) |
|
|
|
|
module_name = "{}.{}".format(anchor_package, proto_name) |
|
|
|
|
if module_name not in sys.modules: |
|
|
|
|
# TODO: The imp module is apparently deprecated. Figure out how to migrate |
|
|
|
|
# over. |
|
|
|
|
module = imp.new_module(module_name) |
|
|
|
|
six.exec_(code, module.__dict__) |
|
|
|
|
sys.modules[module_name] = module |
|
|
|
|
modules.append(module) |
|
|
|
|
else: |
|
|
|
|
modules.append(sys.modules[module_name]) |
|
|
|
|
return tuple(modules) |
|
|
|
|
return os.path.sep.join(components[:-1] + [proto_base_name + suffix]) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@contextlib.contextmanager |
|
|
|
|
def _augmented_syspath(new_paths): |
|
|
|
|
original_sys_path = sys.path |
|
|
|
|
if new_paths is not None: |
|
|
|
|
sys.path = sys.path + new_paths |
|
|
|
|
try: |
|
|
|
|
yield |
|
|
|
|
finally: |
|
|
|
|
sys.path = original_sys_path |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# TODO: Investigate making this even more of a no-op in the case that we have |
|
|
|
|
# truly already imported the module. |
|
|
|
|
def get_protos(protobuf_path, include_paths=None): |
|
|
|
|
original_sys_path = sys.path |
|
|
|
|
if include_paths is not None: |
|
|
|
|
sys.path = sys.path + include_paths |
|
|
|
|
module_name = _proto_file_to_module_name(protobuf_path) |
|
|
|
|
module = importlib.import_module(module_name) |
|
|
|
|
sys.path = original_sys_path |
|
|
|
|
return module |
|
|
|
|
with _augmented_syspath(include_paths): |
|
|
|
|
# TODO: Pull these strings out to module-level constants. |
|
|
|
|
module_name = _proto_file_to_module_name("_pb2", protobuf_path) |
|
|
|
|
module = importlib.import_module(module_name) |
|
|
|
|
return module |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_services(protobuf_path, include_paths=None): |
|
|
|
|
# NOTE: This call to get_protos is a no-op in the case it has already been |
|
|
|
|
# called. |
|
|
|
|
get_protos(protobuf_path, include_paths) |
|
|
|
|
if include_paths is None: |
|
|
|
|
include_paths = sys.path |
|
|
|
|
files = _protoc_compiler.get_services(protobuf_path.encode('ascii'), [include_path.encode('ascii') for include_path in include_paths]) |
|
|
|
|
return _import_modules_from_files(files)[-1] |
|
|
|
|
with _augmented_syspath(include_paths): |
|
|
|
|
module_name = _proto_file_to_module_name("_pb2_grpc", protobuf_path) |
|
|
|
|
module = importlib.import_module(module_name) |
|
|
|
|
return module |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_protos_and_services(protobuf_path, include_paths=None): |
|
|
|
|
return (get_protos(protobuf_path, include_paths=include_paths), |
|
|
|
@ -94,10 +86,10 @@ def get_protos_and_services(protobuf_path, include_paths=None): |
|
|
|
|
|
|
|
|
|
_proto_code_cache = {} |
|
|
|
|
|
|
|
|
|
# TODO: Cache generated code per-process. Check it first to see if it's already |
|
|
|
|
# been generated and, instead, just instantiate using that. |
|
|
|
|
class ProtoLoader(importlib.abc.Loader): |
|
|
|
|
def __init__(self, module_name, protobuf_path, proto_root): |
|
|
|
|
def __init__(self, suffix, code_fn, module_name, protobuf_path, proto_root): |
|
|
|
|
self._suffix = suffix |
|
|
|
|
self._code_fn = code_fn |
|
|
|
|
self._module_name = module_name |
|
|
|
|
self._protobuf_path = protobuf_path |
|
|
|
|
self._proto_root = proto_root |
|
|
|
@ -116,7 +108,7 @@ class ProtoLoader(importlib.abc.Loader): |
|
|
|
|
code = _proto_code_cache[self._module_name] |
|
|
|
|
six.exec_(code, module.__dict__) |
|
|
|
|
else: |
|
|
|
|
files = _protoc_compiler.get_protos(self._protobuf_path.encode('ascii'), [path.encode('ascii') for path in sys.path]) |
|
|
|
|
files = self._code_fn(self._protobuf_path.encode('ascii'), [path.encode('ascii') for path in sys.path]) |
|
|
|
|
for f in files[:-1]: |
|
|
|
|
module_name = self._generated_file_to_module_name(f[0].decode('ascii')) |
|
|
|
|
if module_name not in sys.modules: |
|
|
|
@ -125,9 +117,14 @@ class ProtoLoader(importlib.abc.Loader): |
|
|
|
|
importlib.import_module(module_name) |
|
|
|
|
six.exec_(files[-1][1], module.__dict__) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ProtoFinder(importlib.abc.MetaPathFinder): |
|
|
|
|
def __init__(self, suffix, code_fn): |
|
|
|
|
self._suffix = suffix |
|
|
|
|
self._code_fn = code_fn |
|
|
|
|
|
|
|
|
|
def find_spec(self, fullname, path, target=None): |
|
|
|
|
filepath = _module_name_to_proto_file(fullname) |
|
|
|
|
filepath = _module_name_to_proto_file(self._suffix, fullname) |
|
|
|
|
for search_path in sys.path: |
|
|
|
|
try: |
|
|
|
|
prospective_path = os.path.join(search_path, filepath) |
|
|
|
@ -136,9 +133,9 @@ class ProtoFinder(importlib.abc.MetaPathFinder): |
|
|
|
|
continue |
|
|
|
|
else: |
|
|
|
|
# TODO: Use a stdlib helper function to construct this. |
|
|
|
|
return importlib.machinery.ModuleSpec(fullname, ProtoLoader(fullname, filepath, search_path)) |
|
|
|
|
return importlib.machinery.ModuleSpec(fullname, ProtoLoader(self._suffix, self._code_fn, fullname, filepath, search_path)) |
|
|
|
|
|
|
|
|
|
sys.meta_path.append(ProtoFinder()) |
|
|
|
|
sys.meta_path.extend([ProtoFinder("_pb2", _protoc_compiler.get_protos), ProtoFinder("_pb2_grpc", _protoc_compiler.get_services)]) |
|
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
proto_include = pkg_resources.resource_filename('grpc_tools', '_proto') |
|
|
|
|