Move parsing of protobuf data types and services into separate methods
This commit is contained in:
boukeversteegh
@@ -163,13 +163,104 @@ def generate_code(request, response):
|
||||
for proto_file in output_package_content["files"]:
|
||||
item: DescriptorProto
|
||||
for item, path in traverse(proto_file):
|
||||
data = {"name": item.name, "py_name": pythonize_class_name(item.name)}
|
||||
read_protobuf_type(input_package_name, item, path, proto_file, template_data)
|
||||
|
||||
for i, service in enumerate(proto_file.service):
|
||||
read_protobuf_service(i, input_package_name, proto_file, service, template_data)
|
||||
|
||||
template_data["imports"] = sorted(template_data["imports"])
|
||||
template_data["datetime_imports"] = sorted(template_data["datetime_imports"])
|
||||
template_data["typing_imports"] = sorted(template_data["typing_imports"])
|
||||
|
||||
# Fill response
|
||||
output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
|
||||
output_paths.add(output_path)
|
||||
|
||||
f = response.file.add()
|
||||
f.name = str(output_path)
|
||||
|
||||
# Render and then format the output file.
|
||||
f.content = black.format_str(
|
||||
template.render(description=template_data),
|
||||
mode=black.FileMode(target_versions=set([black.TargetVersion.PY37])),
|
||||
)
|
||||
|
||||
# Make each output directory a package with __init__ file
|
||||
init_files = (
|
||||
set(
|
||||
directory.joinpath("__init__.py")
|
||||
for path in output_paths
|
||||
for directory in path.parents
|
||||
)
|
||||
- output_paths
|
||||
)
|
||||
|
||||
for init_file in init_files:
|
||||
init = response.file.add()
|
||||
init.name = str(init_file)
|
||||
|
||||
for output_package_name in sorted(output_paths.union(init_files)):
|
||||
print(f"Writing {output_package_name}", file=sys.stderr)
|
||||
|
||||
|
||||
def read_protobuf_service(i, input_package_name, proto_file, service, template_data):
|
||||
# print(service, file=sys.stderr)
|
||||
data = {
|
||||
"name": service.name,
|
||||
"py_name": pythonize_class_name(service.name),
|
||||
"comment": get_comment(proto_file, [6, i]),
|
||||
"methods": [],
|
||||
}
|
||||
for j, method in enumerate(service.method):
|
||||
input_message = None
|
||||
input_type = get_type_reference(
|
||||
input_package_name, template_data["imports"], method.input_type
|
||||
).strip('"')
|
||||
for msg in template_data["messages"]:
|
||||
if msg["name"] == input_type:
|
||||
input_message = msg
|
||||
for field in msg["properties"]:
|
||||
if field["zero"] == "None":
|
||||
template_data["typing_imports"].add("Optional")
|
||||
break
|
||||
|
||||
data["methods"].append(
|
||||
{
|
||||
"name": method.name,
|
||||
"py_name": pythonize_method_name(method.name),
|
||||
"comment": get_comment(proto_file, [6, i, 2, j], indent=8),
|
||||
"route": f"/{input_package_name}.{service.name}/{method.name}",
|
||||
"input": get_type_reference(
|
||||
input_package_name, template_data["imports"], method.input_type
|
||||
).strip('"'),
|
||||
"input_message": input_message,
|
||||
"output": get_type_reference(
|
||||
input_package_name,
|
||||
template_data["imports"],
|
||||
method.output_type,
|
||||
unwrap=False,
|
||||
),
|
||||
"client_streaming": method.client_streaming,
|
||||
"server_streaming": method.server_streaming,
|
||||
}
|
||||
)
|
||||
|
||||
if method.client_streaming:
|
||||
template_data["typing_imports"].add("AsyncIterable")
|
||||
template_data["typing_imports"].add("Iterable")
|
||||
template_data["typing_imports"].add("Union")
|
||||
if method.server_streaming:
|
||||
template_data["typing_imports"].add("AsyncIterator")
|
||||
template_data["services"].append(data)
|
||||
|
||||
|
||||
def read_protobuf_type(input_package_name, item, path, proto_file, template_data):
|
||||
data = {"name": item.name, "py_name": pythonize_class_name(item.name)}
|
||||
if isinstance(item, DescriptorProto):
|
||||
# print(item, file=sys.stderr)
|
||||
if item.options.map_entry:
|
||||
# Skip generated map entry messages since we just use dicts
|
||||
continue
|
||||
return
|
||||
|
||||
data.update(
|
||||
{
|
||||
@@ -291,93 +382,6 @@ def generate_code(request, response):
|
||||
|
||||
template_data["enums"].append(data)
|
||||
|
||||
for i, service in enumerate(proto_file.service):
|
||||
# print(service, file=sys.stderr)
|
||||
|
||||
data = {
|
||||
"name": service.name,
|
||||
"py_name": pythonize_class_name(service.name),
|
||||
"comment": get_comment(proto_file, [6, i]),
|
||||
"methods": [],
|
||||
}
|
||||
|
||||
for j, method in enumerate(service.method):
|
||||
input_message = None
|
||||
input_type = get_type_reference(
|
||||
input_package_name, template_data["imports"], method.input_type
|
||||
).strip('"')
|
||||
for msg in template_data["messages"]:
|
||||
if msg["name"] == input_type:
|
||||
input_message = msg
|
||||
for field in msg["properties"]:
|
||||
if field["zero"] == "None":
|
||||
template_data["typing_imports"].add("Optional")
|
||||
break
|
||||
|
||||
data["methods"].append(
|
||||
{
|
||||
"name": method.name,
|
||||
"py_name": pythonize_method_name(method.name),
|
||||
"comment": get_comment(proto_file, [6, i, 2, j], indent=8),
|
||||
"route": f"/{input_package_name}.{service.name}/{method.name}",
|
||||
"input": get_type_reference(
|
||||
input_package_name, template_data["imports"], method.input_type
|
||||
).strip('"'),
|
||||
"input_message": input_message,
|
||||
"output": get_type_reference(
|
||||
input_package_name,
|
||||
template_data["imports"],
|
||||
method.output_type,
|
||||
unwrap=False,
|
||||
),
|
||||
"client_streaming": method.client_streaming,
|
||||
"server_streaming": method.server_streaming,
|
||||
}
|
||||
)
|
||||
|
||||
if method.client_streaming:
|
||||
template_data["typing_imports"].add("AsyncIterable")
|
||||
template_data["typing_imports"].add("Iterable")
|
||||
template_data["typing_imports"].add("Union")
|
||||
if method.server_streaming:
|
||||
template_data["typing_imports"].add("AsyncIterator")
|
||||
|
||||
template_data["services"].append(data)
|
||||
|
||||
template_data["imports"] = sorted(template_data["imports"])
|
||||
template_data["datetime_imports"] = sorted(template_data["datetime_imports"])
|
||||
template_data["typing_imports"] = sorted(template_data["typing_imports"])
|
||||
|
||||
# Fill response
|
||||
output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
|
||||
output_paths.add(output_path)
|
||||
|
||||
f = response.file.add()
|
||||
f.name = str(output_path)
|
||||
|
||||
# Render and then format the output file.
|
||||
f.content = black.format_str(
|
||||
template.render(description=template_data),
|
||||
mode=black.FileMode(target_versions=set([black.TargetVersion.PY37])),
|
||||
)
|
||||
|
||||
# Make each output directory a package with __init__ file
|
||||
init_files = (
|
||||
set(
|
||||
directory.joinpath("__init__.py")
|
||||
for path in output_paths
|
||||
for directory in path.parents
|
||||
)
|
||||
- output_paths
|
||||
)
|
||||
|
||||
for init_file in init_files:
|
||||
init = response.file.add()
|
||||
init.name = str(init_file)
|
||||
|
||||
for output_package_name in sorted(output_paths.union(init_files)):
|
||||
print(f"Writing {output_package_name}", file=sys.stderr)
|
||||
|
||||
|
||||
def main():
|
||||
"""The plugin's main entry point."""
|
||||
|
||||
Reference in New Issue
Block a user