From 21d88ca079fdcb892465b13aadf9df8bd5afa605 Mon Sep 17 00:00:00 2001 From: Nash Aschenbath Date: Sun, 8 Nov 2026 19:24:14 +0800 Subject: [PATCH] Index source locations per compiler request for comment rendering --- src/betterproto/plugin/models.py | 40 +++++++- tests/test_comments.py | 167 +++++++++++++++++++++++++++++++ 2 files changed, 203 insertions(+), 4 deletions(-) create mode 100644 tests/test_comments.py diff --git a/src/betterproto/plugin/models.py b/src/betterproto/plugin/models.py index e330e688..57bbd9b2 100644 --- a/src/betterproto/plugin/models.py +++ b/src/betterproto/plugin/models.py @@ -42,6 +42,7 @@ List, Optional, Set, + Tuple, Type, Union, ) @@ -61,6 +62,7 @@ FieldDescriptorProtoType, FileDescriptorProto, MethodDescriptorProto, + SourceCodeInfoLocation, ) from betterproto.lib.google.protobuf.compiler import CodeGeneratorRequest @@ -149,10 +151,16 @@ def monkey_patch_oneof_index(): def get_comment( - proto_file: "FileDescriptorProto", path: List[int], indent: int = 4 + proto_file: "FileDescriptorProto", + path: List[int], + indent: int = 4, + *, + locations: Optional[Iterable[SourceCodeInfoLocation]] = None, ) -> str: pad = " " * indent - for sci_loc in proto_file.source_code_info.location: + if locations is None: + locations = proto_file.source_code_info.location + for sci_loc in locations: if list(sci_loc.path) == path: all_comments = list(sci_loc.leading_detached_comments) if sci_loc.leading_comments: @@ -224,7 +232,7 @@ def comment(self) -> str: """Crawl the proto source code and retrieve comments for this object. """ - return get_comment( + return self.request.get_comment( proto_file=self.source_file, path=self.path, indent=self.comment_indent ) @@ -233,6 +241,30 @@ def comment(self) -> str: class PluginRequestCompiler: plugin_request_obj: CodeGeneratorRequest output_packages: Dict[str, "OutputTemplate"] = field(default_factory=dict) + _comment_locations: Dict[int, Dict[Tuple[int, ...], SourceCodeInfoLocation]] = ( + field(default_factory=dict, init=False, repr=False) + ) + + def get_comment( + self, proto_file: FileDescriptorProto, path: List[int], indent: int = 4 + ) -> str: + # Source files live for the duration of this request. Index each one once, + # rather than scanning all its locations for every generated declaration. + file_id = id(proto_file) + if file_id not in self._comment_locations: + locations = {} + for location in proto_file.source_code_info.location: + # Match the first location, as the linear search does. + locations.setdefault(tuple(location.path), location) + self._comment_locations[file_id] = locations + + location = self._comment_locations[file_id].get(tuple(path)) + return get_comment( + proto_file, + path, + indent, + locations=() if location is None else (location,), + ) @property def all_messages(self) -> List["MessageCompiler"]: @@ -659,7 +691,7 @@ def __post_init__(self) -> None: entry_proto_value.name, self.proto_obj.name ), value=entry_proto_value.number, - comment=get_comment( + comment=self.request.get_comment( proto_file=self.source_file, path=self.path + [2, entry_number] ), ) diff --git a/tests/test_comments.py b/tests/test_comments.py new file mode 100644 index 00000000..d6710db7 --- /dev/null +++ b/tests/test_comments.py @@ -0,0 +1,167 @@ +import pytest + +from betterproto.lib.google.protobuf import ( + DescriptorProto, + EnumDescriptorProto, + EnumValueDescriptorProto, + FieldDescriptorProto, + FieldDescriptorProtoType, + FileDescriptorProto, + MethodDescriptorProto, + ServiceDescriptorProto, + SourceCodeInfo, + SourceCodeInfoLocation, +) +from betterproto.lib.google.protobuf.compiler import CodeGeneratorRequest +from betterproto.plugin.models import PluginRequestCompiler, get_comment +from betterproto.plugin.parser import generate_code + + +@pytest.mark.parametrize("indent", [0, 4, 8]) +def test_indexed_comments_preserve_formatting(indent): + locations = [ + SourceCodeInfoLocation( + path=[4, i], + leading_detached_comments=detached, + leading_comments=leading, + trailing_comments=trailing, + ) + for i, (detached, leading, trailing) in enumerate( + [ + ([], "", ""), + ([], " A short comment.", ""), + ([], "\n\n Leading\n\n\n lines.\n", ""), + ( + [" Detached one.\n", " Detached two.\n"], + " Leading.\n", + " Trailing.\n", + ), + ([], " A long comment " * 10, ""), + ([], "\tUnicode: \u03bb\u4e2d\u6587", ""), + ] + ) + ] + # A repeated path must keep the first location, even if it has no comment. + locations.append(SourceCodeInfoLocation(path=[4, 0], leading_comments=" Later.")) + source = FileDescriptorProto(source_code_info=SourceCodeInfo(location=locations)) + compiler = PluginRequestCompiler(CodeGeneratorRequest(proto_file=[source])) + + for path in [[4, i] for i in range(6)] + [[4, 99], [], [4]]: + assert compiler.get_comment(source, path, indent) == get_comment( + source, path, indent + ) + + +def test_comment_index_is_scoped_to_source_file_and_request(): + first = FileDescriptorProto( + name="first.proto", + source_code_info=SourceCodeInfo( + location=[SourceCodeInfoLocation(path=[4, 0], leading_comments=" First.")] + ), + ) + second = FileDescriptorProto( + name="second.proto", + source_code_info=SourceCodeInfo( + location=[SourceCodeInfoLocation(path=[4, 0], leading_comments=" Second.")] + ), + ) + request = CodeGeneratorRequest(proto_file=[first, second]) + compiler = PluginRequestCompiler(request) + assert compiler.get_comment(first, [4, 0]) == ' """First."""' + assert compiler.get_comment(second, [4, 0]) == ' """Second."""' + + first.source_code_info.location[0].path = [4, 1] + next_compiler = PluginRequestCompiler(request) + assert next_compiler.get_comment(first, [4, 0]) == "" + assert next_compiler.get_comment(first, [4, 1]) == ' """First."""' + + +def test_generated_comments_do_not_repeatedly_scan_locations(monkeypatch): + field_count = 100 + sources = [] + for file_index in range(2): + locations = [ + SourceCodeInfoLocation( + path=[4, 0], leading_comments=f" Message {file_index}." + ) + ] + fields = [] + for i in range(field_count): + fields.append( + FieldDescriptorProto( + name=f"field_{i}", + number=i + 1, + type=FieldDescriptorProtoType.TYPE_STRING, + ) + ) + locations.append( + SourceCodeInfoLocation( + path=[4, 0, 2, i], leading_comments=f" Field {file_index}/{i}." + ) + ) + locations += [ + SourceCodeInfoLocation( + path=[5, 0], leading_comments=f" Enum {file_index}." + ), + SourceCodeInfoLocation( + path=[5, 0, 2, 0], leading_comments=f" Entry {file_index}." + ), + SourceCodeInfoLocation( + path=[6, 0], leading_comments=f" Service {file_index}." + ), + SourceCodeInfoLocation( + path=[6, 0, 2, 0], leading_comments=f" Method {file_index}." + ), + ] + sources.append( + FileDescriptorProto( + name=f"file_{file_index}.proto", + package="comments", + message_type=[ + DescriptorProto(name=f"Message{file_index}", field=fields) + ], + enum_type=[ + EnumDescriptorProto( + name=f"Enum{file_index}", + value=[EnumValueDescriptorProto(name="ZERO", number=0)], + ) + ], + service=[ + ServiceDescriptorProto( + name=f"Service{file_index}", + method=[ + MethodDescriptorProto( + name="Call", + input_type=f".comments.Message{file_index}", + output_type=f".comments.Message{file_index}", + ) + ], + ) + ], + source_code_info=SourceCodeInfo(location=locations), + ) + ) + + path_reads = 0 + original_getattribute = SourceCodeInfoLocation.__getattribute__ + + def count_path_reads(self, name): + nonlocal path_reads + if name == "path": + path_reads += 1 + return original_getattribute(self, name) + + monkeypatch.setattr(SourceCodeInfoLocation, "__getattribute__", count_path_reads) + response = generate_code(CodeGeneratorRequest(proto_file=sources)) + content = "\n".join(file.content for file in response.file) + for file_index in range(2): + for kind in ("Message", "Enum", "Entry", "Service", "Method"): + assert f"{kind} {file_index}." in content + for i in range(field_count): + assert f"Field {file_index}/{i}." in content + + # Each source location may be indexed and checked a few times, but the + # number of reads must grow linearly with the number of declarations. + assert path_reads < 10 * sum( + len(source.source_code_info.location) for source in sources + )