fix: register descriptor dependencies before AddSerializedFile (#150)

betterproto2/docs/descriptors.md

@@ -4,6 +4,8 @@ Google's protoc plugin for Python generated DESCRIPTOR fields that enable reflec
 
 By default, betterproto2 doesn't generate these as it introduces a dependency on `protobuf`. If you're okay with this dependency and want to generate DESCRIPTORs, use the compiler option `python_betterproto2_opt=google_protobuf_descriptors`.
 
+Generated modules register each file's `FileDescriptor` with a shared `DescriptorPool` at import time. Files that import another proto (including `google/protobuf/descriptor.proto` for custom options) first import that proto's generated package so its descriptors are already in the pool. Package-level import cycles are not handled; the proto file graph is expected to be a DAG.
+
 
 ## grpclib Reflection
 

betterproto2/tests/grpc/test_descriptor_pool_deps.py

@@ -0,0 +1,25 @@
+from tests.util import requires_protobuf  # noqa: F401
+
+
+def test_descriptor_options_import(requires_protobuf):
+    """Custom options that import google/protobuf/descriptor.proto should import.
+
+    Reproduces https://github.com/betterproto/python-betterproto2/issues/150.
+    """
+    from tests.outputs.descriptor_options_descriptors.descriptor_options import MyMessage
+
+    assert MyMessage.DESCRIPTOR.full_name == "descriptor_options.MyMessage"
+
+
+def test_descriptor_cross_package_import(requires_protobuf):
+    """A package whose proto imports another package must load that package first."""
+    from tests.outputs.descriptor_cross_package_descriptors.descriptor_cross_package.use import UseMsg
+
+    assert UseMsg.DESCRIPTOR.full_name == "descriptor_cross_package.use.UseMsg"
+
+
+def test_descriptor_same_package_import(requires_protobuf):
+    """Files in one package must be registered in dependency order."""
+    from tests.outputs.descriptor_same_package_descriptors.descriptor_same_package import Depends
+
+    assert Depends.DESCRIPTOR.full_name == "descriptor_same_package.Depends"

betterproto2_compiler/src/betterproto2_compiler/plugin/models.py

@@ -32,7 +32,11 @@
 from betterproto2 import unwrap
 
 from betterproto2_compiler import casing
-from betterproto2_compiler.compile.importing import get_type_reference, parse_source_type_name
+from betterproto2_compiler.compile.importing import (
+    get_symbol_reference,
+    get_type_reference,
+    parse_source_type_name,
+)
 from betterproto2_compiler.compile.naming import (
     pythonize_class_name,
     pythonize_field_name,
@@ -217,18 +221,81 @@ def input_filenames(self) -> list[str]:
     def get_descriptor_name(self, source_file: FileDescriptorProto):
         return f"{source_file.name.replace('/', '_').replace('.', '_').upper()}_DESCRIPTOR"
 
+    def _ordered_input_files(self) -> list[FileDescriptorProto]:
+        """Return input files in proto dependency order.
+
+        Files from other packages are skipped; those are loaded via ``_descriptor_dependency_imports`` instead.
+        """
+        by_name = {proto_file.name: proto_file for proto_file in self.input_files}
+        ordered: list[FileDescriptorProto] = []
+        visiting: set[str] = set()
+        visited: set[str] = set()
+
+        def _visit(name: str) -> None:
+            if name in visited or name not in by_name:
+                return
+            if name in visiting:
+                return
+            visiting.add(name)
+            for dep_name in by_name[name].dependency:
+                _visit(dep_name)
+            visiting.remove(name)
+            visited.add(name)
+            ordered.append(by_name[name])
+
+        for proto_file in self.input_files:
+            _visit(proto_file.name)
+
+        return ordered
+
+    def _descriptor_dependency_imports(self) -> list[str]:
+        """Relative imports that register foreign-package descriptors first."""
+        file_to_package = {
+            proto_file.name: package_name
+            for package_name, package in self.parent_request.output_packages.items()
+            for proto_file in package.input_files
+        }
+
+        needed: set[str] = set()
+        for proto_file in self.input_files:
+            for dep_name in proto_file.dependency:
+                dep_package = file_to_package.get(dep_name)
+                if dep_package is not None and dep_package != self.package:
+                    needed.add(dep_package)
+
+        imports: set[str] = set()
+        aliases: list[str] = []
+        for dep_package in sorted(needed):
+            ref, _ = get_symbol_reference(
+                package=self.package,
+                imports=imports,
+                source_package=dep_package,
+                symbol="_prereq",
+            )
+            alias = ref.rsplit(".", 1)[0]
+            aliases.append(alias)
+
+        lines: list[str] = []
+        for imp in sorted(imports):
+            lines.append(imp)
+        for alias in aliases:
+            # Keep the import live so ruff F401 does not strip the side-effect load.
+            lines.append(f"_ = {alias}")
+        return lines
+
     @property
     def descriptors(self):
         """Google protobuf library descriptors.
 
         Returns
         -------
         str
-            A list of pool registrations for proto descriptors.
+            Prerequisite package imports followed by pool registrations.
         """
         descriptors: list[str] = []
+        descriptors.extend(self._descriptor_dependency_imports())
 
-        for f in self.input_files:
+        for f in self._ordered_input_files():
             # Remove the source_code_info field since it is not needed at runtime.
             source_code_info: SourceCodeInfo | None = f.source_code_info
             f.source_code_info = None

betterproto2_compiler/src/betterproto2_compiler/templates/template.py.j2

@@ -133,7 +133,7 @@ default_message_pool.register_message("{{ output_file.package }}", "{{ message.p
 {% endfor %}
 
 {% if output_file.settings.google_protobuf_descriptors %}
-{# Add descriptors to Google protobuf's default pool to be more drop-in compatible with other libraries. #}
+{# Register dependency descriptors to Google protobuf's default pool, then this file's FileDescriptorProtos. #}
 {{ output_file.descriptors }}
 {% endif %}
 

betterproto2_compiler/tests/generate.py

@@ -77,6 +77,9 @@ async def main_async(outputs_dir: str):
         generate_test("casing", semaphore, outputs_dir),
         generate_test("compiler_lib", semaphore, outputs_dir),
         generate_test("conformance", semaphore, outputs_dir),
+        generate_test("descriptor_cross_package", semaphore, outputs_dir, descriptors=True),
+        generate_test("descriptor_options", semaphore, outputs_dir, descriptors=True),
+        generate_test("descriptor_same_package", semaphore, outputs_dir, descriptors=True),
         generate_test("deprecated", semaphore, outputs_dir, reference=True),
         generate_test("deprecated", semaphore, outputs_dir, client_generation="async"),
         generate_test("documentation", semaphore, outputs_dir, client_generation="async"),

betterproto2_compiler/tests/inputs/descriptor_cross_package/dep.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package descriptor_cross_package.dep;
+
+message DepMsg {
+  string name = 1;
+}

betterproto2_compiler/tests/inputs/descriptor_cross_package/use.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package descriptor_cross_package.use;
+
+import "dep.proto";
+
+message UseMsg {
+  descriptor_cross_package.dep.DepMsg dep = 1;
+}

betterproto2_compiler/tests/inputs/descriptor_options/example.proto

@@ -0,0 +1,14 @@
+syntax = "proto3";
+
+package descriptor_options;
+
+import "google/protobuf/descriptor.proto";
+
+extend google.protobuf.MessageOptions {
+  string example_option = 50000;
+}
+
+message MyMessage {
+  option (example_option) = "bar";
+  string baz = 1;
+}

betterproto2_compiler/tests/inputs/descriptor_same_package/a_depends.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package descriptor_same_package;
+
+import "z_base.proto";
+
+message Depends {
+  Base base = 1;
+}

betterproto2_compiler/tests/inputs/descriptor_same_package/z_base.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package descriptor_same_package;
+
+message Base {
+  string name = 1;
+}

betterproto2_compiler/tests/util.py

@@ -24,7 +24,7 @@ async def protoc(
         "grpc.tools.protoc",
         f"--proto_path={resolved_path.as_posix()}",
         f"--{python_out_option}={resolved_output_dir.as_posix()}",
-        *[p.as_posix() for p in resolved_path.glob("*.proto")],
+        *[p.as_posix() for p in sorted(resolved_path.glob("*.proto"))],
     ]
 
     if not reference: