Make plugin use betterproto generated classes internally
This means the betterproto plugin no longer needs to depend durectly on protobuf. This requires a small runtime hack to monkey patch some google types to get around the fact that the compiler uses proto2, but betterproto expects proto3. Also: - regenerate google.protobuf package - fix a regex bug in the logic for determining whether to use a google wrapper type. - fix a bug causing comments to get mixed up when multiple proto files generate code into a single python module
This commit is contained in:
@@ -1,28 +1,19 @@
|
||||
from betterproto.lib.google.protobuf import (
|
||||
DescriptorProto,
|
||||
EnumDescriptorProto,
|
||||
FieldDescriptorProto,
|
||||
FileDescriptorProto,
|
||||
ServiceDescriptorProto,
|
||||
)
|
||||
from betterproto.lib.google.protobuf.compiler import (
|
||||
CodeGeneratorRequest,
|
||||
CodeGeneratorResponse,
|
||||
CodeGeneratorResponseFile,
|
||||
)
|
||||
import itertools
|
||||
import pathlib
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Iterator, List, Tuple, Union, Set
|
||||
|
||||
try:
|
||||
# betterproto[compiler] specific dependencies
|
||||
from google.protobuf.compiler import plugin_pb2 as plugin
|
||||
from google.protobuf.descriptor_pb2 import (
|
||||
DescriptorProto,
|
||||
EnumDescriptorProto,
|
||||
FieldDescriptorProto,
|
||||
ServiceDescriptorProto,
|
||||
)
|
||||
except ImportError as err:
|
||||
print(
|
||||
"\033[31m"
|
||||
f"Unable to import `{err.name}` from betterproto plugin! "
|
||||
"Please ensure that you've installed betterproto as "
|
||||
'`pip install "betterproto[compiler]"` so that compiler dependencies '
|
||||
"are included."
|
||||
"\033[0m"
|
||||
)
|
||||
raise SystemExit(1)
|
||||
|
||||
from typing import Iterator, List, Tuple, TYPE_CHECKING, Union
|
||||
from .compiler import outputfile_compiler
|
||||
from .models import (
|
||||
EnumDefinitionCompiler,
|
||||
@@ -70,7 +61,7 @@ def traverse(
|
||||
|
||||
|
||||
def generate_code(
|
||||
request: plugin.CodeGeneratorRequest, response: plugin.CodeGeneratorResponse
|
||||
request: CodeGeneratorRequest, response: CodeGeneratorResponse
|
||||
) -> None:
|
||||
plugin_options = request.parameter.split(",") if request.parameter else []
|
||||
|
||||
@@ -100,7 +91,12 @@ def generate_code(
|
||||
for output_package_name, output_package in request_data.output_packages.items():
|
||||
for proto_input_file in output_package.input_files:
|
||||
for item, path in traverse(proto_input_file):
|
||||
read_protobuf_type(item=item, path=path, output_package=output_package)
|
||||
read_protobuf_type(
|
||||
source_file=proto_input_file,
|
||||
item=item,
|
||||
path=path,
|
||||
output_package=output_package,
|
||||
)
|
||||
|
||||
# Read Services
|
||||
for output_package_name, output_package in request_data.output_packages.items():
|
||||
@@ -116,11 +112,13 @@ def generate_code(
|
||||
output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
|
||||
output_paths.add(output_path)
|
||||
|
||||
f: response.File = response.file.add()
|
||||
f.name = str(output_path)
|
||||
|
||||
# Render and then format the output file
|
||||
f.content = outputfile_compiler(output_file=output_package)
|
||||
response.file.append(
|
||||
CodeGeneratorResponseFile(
|
||||
name=str(output_path),
|
||||
# Render and then format the output file
|
||||
content=outputfile_compiler(output_file=output_package),
|
||||
)
|
||||
)
|
||||
|
||||
# Make each output directory a package with __init__ file
|
||||
init_files = {
|
||||
@@ -130,38 +128,53 @@ def generate_code(
|
||||
} - output_paths
|
||||
|
||||
for init_file in init_files:
|
||||
init = response.file.add()
|
||||
init.name = str(init_file)
|
||||
response.file.append(CodeGeneratorResponseFile(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_type(
|
||||
item: DescriptorProto, path: List[int], output_package: OutputTemplate
|
||||
item: DescriptorProto,
|
||||
path: List[int],
|
||||
source_file: "FileDescriptorProto",
|
||||
output_package: OutputTemplate,
|
||||
) -> None:
|
||||
if isinstance(item, DescriptorProto):
|
||||
if item.options.map_entry:
|
||||
# Skip generated map entry messages since we just use dicts
|
||||
return
|
||||
# Process Message
|
||||
message_data = MessageCompiler(parent=output_package, proto_obj=item, path=path)
|
||||
message_data = MessageCompiler(
|
||||
source_file=source_file, parent=output_package, proto_obj=item, path=path
|
||||
)
|
||||
for index, field in enumerate(item.field):
|
||||
if is_map(field, item):
|
||||
MapEntryCompiler(
|
||||
parent=message_data, proto_obj=field, path=path + [2, index]
|
||||
source_file=source_file,
|
||||
parent=message_data,
|
||||
proto_obj=field,
|
||||
path=path + [2, index],
|
||||
)
|
||||
elif is_oneof(field):
|
||||
OneOfFieldCompiler(
|
||||
parent=message_data, proto_obj=field, path=path + [2, index]
|
||||
source_file=source_file,
|
||||
parent=message_data,
|
||||
proto_obj=field,
|
||||
path=path + [2, index],
|
||||
)
|
||||
else:
|
||||
FieldCompiler(
|
||||
parent=message_data, proto_obj=field, path=path + [2, index]
|
||||
source_file=source_file,
|
||||
parent=message_data,
|
||||
proto_obj=field,
|
||||
path=path + [2, index],
|
||||
)
|
||||
elif isinstance(item, EnumDescriptorProto):
|
||||
# Enum
|
||||
EnumDefinitionCompiler(parent=output_package, proto_obj=item, path=path)
|
||||
EnumDefinitionCompiler(
|
||||
source_file=source_file, parent=output_package, proto_obj=item, path=path
|
||||
)
|
||||
|
||||
|
||||
def read_protobuf_service(
|
||||
|
||||
Reference in New Issue
Block a user