Support services too

pull/21504/head
Richard Belleville 5 years ago
parent 595807e0c5
commit ecab62dc31
  1. 83
      tools/distrib/python/grpcio_tools/grpc_tools/protoc.py
  2. 45
      tools/distrib/python/grpcio_tools/grpc_tools/protoc_test.py

@ -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')

@ -100,42 +100,31 @@ def _test_syntax_errors():
class ProtocTest(unittest.TestCase):
# def test_import_protos(self):
# _run_in_subprocess(_test_import_protos)
def test_import_protos(self):
_run_in_subprocess(_test_import_protos)
# def test_import_services(self):
# _run_in_subprocess(_test_import_services)
def test_import_services(self):
_run_in_subprocess(_test_import_services)
# def test_import_implicit_include_path(self):
# _run_in_subprocess(_test_import_implicit_include)
def test_import_implicit_include_path(self):
_run_in_subprocess(_test_import_implicit_include)
# def test_import_services_without_protos(self):
# _run_in_subprocess(_test_import_services_without_protos)
def test_import_services_without_protos(self):
_run_in_subprocess(_test_import_services_without_protos)
# def test_proto_module_imported_once(self):
# _run_in_subprocess(_test_proto_module_imported_once)
def test_proto_module_imported_once(self):
_run_in_subprocess(_test_proto_module_imported_once)
# def test_static_dynamic_combo(self):
# _run_in_subprocess(_test_static_dynamic_combo)
def test_static_dynamic_combo(self):
_run_in_subprocess(_test_static_dynamic_combo)
# def test_combined_import(self):
# _run_in_subprocess(_test_combined_import)
def test_combined_import(self):
_run_in_subprocess(_test_combined_import)
# def test_syntax_errors(self):
# _run_in_subprocess(_test_syntax_errors)
def test_syntax_errors(self):
_run_in_subprocess(_test_syntax_errors)
# # TODO: Write test to ensure the right module loader is used.
# def test_importlib_protos(self):
# import sys
# import grpc_tools.protoc
# from grpc_tools import simple_pb2
# self.assertIsNotNone(simple_pb2.SimpleMessage)
def test_importlib_protos_wrapper(self):
from grpc_tools import protoc
proto_path = "tools/distrib/python/grpcio_tools/"
protos = protoc.get_protos("grpc_tools/simple.proto", [proto_path])
assert protos.SimpleMessage is not None
# TODO: Write test to ensure the right module loader is used.
if __name__ == '__main__':
unittest.main()

Loading…
Cancel
Save