def load_descriptor_set(
proto_file: str, proto_path: Optional[List[str]] = None
) -> FileDescriptorSet:
wkt_protos = pkg_resources.resource_filename("grpc_tools", "_proto")
validate_protos = str(
Path(
pkg_resources.resource_filename("validate", "validate.proto")
).parent.parent
)
harness_protos = str(Path(__file__).parent.parent.parent.absolute())
with tempfile.NamedTemporaryFile() as f:
args = [
"grpc_tools.protoc",
"--include_imports",
f"--proto_path={wkt_protos}",
f"--proto_path={validate_protos}",
f"--proto_path={harness_protos}",
f"--proto_path={Path(proto_file).parent}",
f"--descriptor_set_out={f.name}",
]
if proto_path:
args.extend(f"--proto_path={p}" for p in proto_path)
args.append(proto_file)
result = protoc.main(args)
if result != 0:
raise Exception("Failed to call protoc")
content = f.read()
return FileDescriptorSet.FromString(content)
def parse(self, descfilename, allmeta):
"""解析传入的desc文件，转换成ProtoMeta信息
Args:
desc_file_path: protoc --descriptor_set_out --include_source_info
编译生成的proto描述文件
Returns:
A dict mapping message name and ProtoMeta
"""
with open(descfilename, "rb") as f:
desc = FileDescriptorSet.FromString(f.read())
for onefile in desc.file:
allmeta.add(self.parse_file(onefile))
# for proto_message_meta in proto_meta_mgr.metas:
# if proto_message_meta.name in self._reference_message_name:
# proto_message_meta.is_reference_message = True
return allmeta
def load_proto_msgs(proto_path, ret_source_info=False):
# List imports that we need to specify to protoc for the necessary *_pb2.py to be generated
proto_dir = Path(proto_path).parent
arg_proto_path = proto_dir
arg_proto_files = []
to_import = [str(proto_path)]
while to_import:
next_import = to_import.pop()
while not exists(arg_proto_path / next_import) and \
str(arg_proto_path.parent).startswith(str(BASE_PATH)):
arg_proto_path = arg_proto_path.parent
next_import = str(arg_proto_path / next_import)
if next_import not in arg_proto_files:
arg_proto_files.insert(0, next_import)
with open(next_import) as fd:
for prior_import in reversed(
findall('import(?:\s*weak|public)?\s*"(.+?)"\s*;',
fd.read())):
to_import.append(prior_import)
# Execute protoc and import the actual module from a tmp
with TemporaryDirectory() as arg_python_out:
args = [
'protoc',
'--proto_path=%s' % arg_proto_path,
'--python_out=' + arg_python_out, *arg_proto_files
]
if ret_source_info:
args += [
'-o%s' % (Path(arg_python_out) / 'desc_info'),
'--include_source_info', '--include_imports'
]
cmd = run(args, stderr=PIPE, encoding='utf8')
if cmd.returncode:
raise ValueError(cmd.stderr)
if ret_source_info:
with open(Path(arg_python_out) / 'desc_info', 'rb') as fd:
yield FileDescriptorSet.FromString(fd.read()), arg_proto_path
return
# Do actual import
module_name = str(proto_dir).replace(str(arg_proto_path),
'').strip('/\\').replace(
'/', '.')
if module_name:
module_name += '.'
module_name += Path(proto_path).stem.replace('-', '_') + '_pb2'
PATH.append(arg_python_out)
module = import_module(module_name)
PATH.remove(arg_python_out)
# Recursively iterate over class members to list Protobuf messages
yield from iterate_proto_msg(module, '')
def read_desc(src):
with open(src, "rb") as f:
return FileDescriptorSet.FromString(f.read())