Add support for nested definitions in Python mcap-protobuf-support

Signed-off-by: methylDragon <[email protected]>

python/mcap-protobuf-support/mcap_protobuf/decoder.py

@@ -4,7 +4,7 @@
 
 from google.protobuf.descriptor_pb2 import FileDescriptorProto, FileDescriptorSet
 from google.protobuf.descriptor_pool import DescriptorPool
-from google.protobuf.message_factory import GetMessageClassesForFiles
+from google.protobuf.message_factory import GetMessageClassesForFiles, GetMessageClass
 
 from mcap.decoder import DecoderFactory as McapDecoderFactory
 from mcap.exceptions import McapError
@@ -59,6 +59,26 @@ def _add(fd: FileDescriptorProto):
 
             messages = GetMessageClassesForFiles([fd.name for fd in fds.file], pool)
 
+            # GetMessageClassesForFiles only fetches top-level message definitions in the file.
+            #
+            # We must traverse the top-level message definitions to recursively find and append any
+            # nested message definitions.
+            nested_messages = {}
+            for message_class in messages.values():
+                nested_descriptions = list(
+                    message_class.DESCRIPTOR.nested_types_by_name.values()
+                )
+                while nested_descriptions:
+                    nested_desc = nested_descriptions.pop()
+                    nested_messages[nested_desc.full_name] = GetMessageClass(
+                        nested_desc
+                    )
+                    nested_descriptions.extend(
+                        nested_desc.nested_types_by_name.values()
+                    )
+
+            messages.update(nested_messages)
+
             for name, klass in messages.items():
                 if name == schema.name:
                     self._types[schema.id] = klass

python/mcap-protobuf-support/tests/generate.py

@@ -13,6 +13,7 @@
 from test_proto.complex_message_pb2 import ComplexMessage  # noqa: #402
 from test_proto.intermediate_message_1_pb2 import IntermediateMessage1  # noqa: #402
 from test_proto.intermediate_message_2_pb2 import IntermediateMessage2  # noqa: #402
+from test_proto.nested_type_message_pb2 import NestedTypeMessage  # noqa: #402
 from test_proto.simple_message_pb2 import SimpleMessage  # noqa: #402
 
 
@@ -65,6 +66,7 @@ def generate_sample_data(output: IO[Any]):
                 log_time=i * 1000,
                 publish_time=i * 1000,
             )
+
             complex_message = ComplexMessage(
                 intermediate1=IntermediateMessage1(
                     simple=SimpleMessage(data=f"Field A {i}")
@@ -79,4 +81,22 @@ def generate_sample_data(output: IO[Any]):
                 log_time=i * 1000,
                 publish_time=i * 1000,
             )
+
+            nested_type_message = NestedTypeMessage(
+                nested_type_message_1=NestedTypeMessage.NestedType1(
+                    doubly_nested_type_message=NestedTypeMessage.NestedType1.DoublyNestedType(
+                        data=f"Field C {i}"
+                    )
+                ),
+                nested_type_message_2=NestedTypeMessage.NestedType2(
+                    data=f"Field D {i}"
+                ),
+            )
+            writer.write_message(
+                topic="/nested_type_message",
+                message=nested_type_message,
+                log_time=i * 1000,
+                publish_time=i * 1000,
+            )
+
     output.seek(0)

python/mcap-protobuf-support/tests/proto/test_proto/nested_type_message.proto

@@ -0,0 +1,21 @@
+syntax = "proto3";
+
+
+package test_proto;
+
+message NestedTypeMessage {
+  message NestedType1 {
+    message DoublyNestedType {
+      string data = 1;
+    }
+
+    DoublyNestedType doubly_nested_type_message = 1;
+  }
+
+  message NestedType2 {
+    string data = 1;
+  }
+
+  NestedType1 nested_type_message_1 = 1;
+  NestedType2 nested_type_message_2 = 2;
+}

python/mcap-protobuf-support/tests/test_decoder.py

@@ -23,12 +23,17 @@ def test_protobuf_decoder():
         if channel.topic == "/complex_message":
             assert proto_msg.intermediate1.simple.data.startswith("Field A")
             assert proto_msg.intermediate2.simple.data.startswith("Field B")
+        elif channel.topic == "/nested_type_message":
+            assert proto_msg.nested_type_message_1.doubly_nested_type_message.data.startswith(
+                "Field C"
+            )
+            assert proto_msg.nested_type_message_2.data.startswith("Field D")
         elif channel.topic == "/simple_message":
             assert proto_msg.data.startswith("Hello MCAP protobuf world")
         else:
             raise AssertionError(f"unrecognized topic {channel.topic}")
 
-    assert count == 20
+    assert count == 30
 
 
 def test_with_disordered_file_descriptors():