feat: Convert Between MCAP and Bag (#125)

Author: siliconlad

benchmarks/writing/test_mcap_writer.py

@@ -83,7 +83,7 @@ def _generate_pybag_odometries(count: int = 1000, seed: int = 0) -> list[Odometr
 
     def _write_with_pybag(path: Path, messages: Iterable) -> None:
         writer = McapFileWriter.open(path)
-        writer.add_channel("/odom", Odometry)
+        writer.add_channel("/odom", schema=Odometry)
         for i, msg in enumerate(messages):
             timestamp = int(i * 1_500_000_000)
             writer.write_message("/odom", timestamp, msg)

src/pybag/bag_writer.py

@@ -20,8 +20,12 @@
 from pybag.encoding.rosmsg import RosMsgEncoder
 from pybag.io.raw_writer import BaseWriter, BytesWriter, FileWriter
 from pybag.schema.ros1_compiler import compile_ros1_serializer
-from pybag.schema.ros1msg import Ros1MsgSchemaEncoder, compute_md5sum
-from pybag.types import Message
+from pybag.schema.ros1msg import (
+    Ros1MsgSchemaDecoder,
+    Ros1MsgSchemaEncoder,
+    compute_md5sum
+)
+from pybag.types import Message, SchemaText
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +46,7 @@ def __init__(
         writer: BaseWriter,
         *,
         compression: Literal['none', 'bz2'] = 'none',
-        chunk_size: int = 1024 * 1024,  # 1MB default chunk size
+        chunk_size: int | None = None,
     ):
         """Initialize the bag writer.
 
@@ -53,11 +57,12 @@ def __init__(
         """
         self._writer = writer
         self._compression = compression
-        self._chunk_size = chunk_size
+        self._chunk_size = chunk_size or (1024 * 1024)  # 1MB
         self._record_writer = BagRecordWriter(writer)
 
-        # Schema encoder
+        # Schema encoder and decoder
         self._schema_encoder = Ros1MsgSchemaEncoder()
+        self._schema_decoder = Ros1MsgSchemaDecoder()
 
         # Tracking state
         self._next_conn_id = 0
@@ -66,13 +71,17 @@ def __init__(
         self._message_types: dict[type[Message], tuple[str, str]] = {}  # type -> (msg_def, md5sum)
         self._serializers: dict[type[Message], Callable[[Any, Any], None]] = {}
 
+        # Pre-compiled serializers for topics with explicit schemas
+        # Maps topic -> compiled serializer function
+        self._topic_serializers: dict[str, Callable[[Any, Any], None]] = {}
+
         # Current chunk state
         self._chunk_buffer = BytesWriter()
         self._chunk_record_writer = BagRecordWriter(self._chunk_buffer)
         self._chunk_start_time: int | None = None
         self._chunk_end_time: int | None = None
         self._chunk_message_counts: dict[int, int] = {}
-        # Index entries for current chunk: conn_id -> [(time_sec, time_nsec, offset)]
+        # Index entries for current chunk: conn_id -> [(time, offset)]
         self._chunk_index_entries: dict[int, list[tuple[int, int]]] = {}
 
         # Chunk info records (for summary)
@@ -87,7 +96,7 @@ def open(
         file_path: str | Path,
         *,
         compression: Literal['none', 'bz2'] = 'none',
-        chunk_size: int = 1024 * 1024,
+        chunk_size: int | None = None,
     ) -> "BagFileWriter":
         """Create a writer for a file.
 
@@ -123,11 +132,7 @@ def _write_header(self) -> None:
         self._header_pos = self._record_writer.tell()
         # Write placeholder header with zeros (will be updated on close)
         self._record_writer.write_bag_header(
-            BagHeaderRecord(
-                index_pos=0,
-                conn_count=0,
-                chunk_count=0,
-            ),
+            BagHeaderRecord(index_pos=0, conn_count=0, chunk_count=0),
         )
 
     def _get_message_info(self, message_type: type[Message]) -> tuple[str, str]:
@@ -161,12 +166,19 @@ def _get_serializer(self, message_type: type[Message]) -> Callable[[Any, Any], N
             self._serializers[message_type] = compile_ros1_serializer(schema, sub_schemas)
         return self._serializers[message_type]
 
-    def add_connection(self, topic: str, message_type: type[Message]) -> int:
+    def add_connection(
+        self,
+        topic: str,
+        *,
+        schema: SchemaText | type[Message] | Message,
+    ) -> int:
         """Add a connection (topic) to the bag file.
 
         Args:
             topic: The topic name.
-            message_type: The message type class.
+            schema: A SchemaText object containing the message type name and
+                   schema definition text, or a message class/instance to
+                   generate the schema from.
 
         Returns:
             The connection ID.
@@ -175,22 +187,38 @@ def add_connection(self, topic: str, message_type: type[Message]) -> int:
         if topic in self._topics:
             return self._topics[topic]
 
+        # Convert message class or instance to SchemaText
+        if isinstance(schema, type) and hasattr(schema, '__msg_name__'):
+            schema = SchemaText(
+               name=schema.__msg_name__,
+               text=self._schema_encoder.encode(schema).decode('utf-8'),
+            )
+        elif isinstance(schema, Message):
+            schema_type = type(schema)
+            schema = SchemaText(
+               name=schema_type.__msg_name__,
+               text=self._schema_encoder.encode(schema_type).decode('utf-8'),
+            )
+
         conn_id = self._next_conn_id
         self._next_conn_id += 1
 
-        msg_def, md5sum = self._get_message_info(message_type)
-        msg_type = message_type.__msg_name__
+        # Use provided schema text directly
+        msg_def = schema.text
+        msg_type_name = schema.name
+        md5sum = compute_md5sum(msg_def, msg_type_name)
 
         # Build the connection data (connection header fields)
         data_buffer = BytesWriter()
         # Two topic fields exist (in the record and connection headers).
         # This is because messages can be written to the bag file on a topic different
         # from where they were originally published
         data_buffer.write(self._encode_header_field('topic', topic.encode('utf-8')))
-        data_buffer.write(self._encode_header_field('type', msg_type.encode('utf-8')))
+        data_buffer.write(self._encode_header_field('type', msg_type_name.encode('utf-8')))
         data_buffer.write(self._encode_header_field('md5sum', md5sum.encode('ascii')))
         data_buffer.write(self._encode_header_field('message_definition', msg_def.encode('utf-8')))
 
+        # TODO: Add checks to see if previous topic exists
         connection = ConnectionRecord(
             conn=conn_id,
             topic=topic,
@@ -202,6 +230,12 @@ def add_connection(self, topic: str, message_type: type[Message]) -> int:
         # Write connection record to current chunk
         self._chunk_record_writer.write_connection(connection)
 
+        # If explicit schema was provided, compile and store a serializer for this topic
+        # This allows us to serialize messages without relying on type annotations
+        parsed_schema, sub_schemas = self._schema_decoder.parse_schema(connection)
+        serializer = compile_ros1_serializer(parsed_schema, sub_schemas)
+        self._topic_serializers[topic] = serializer
+
         return conn_id
 
     def write_message(
@@ -212,15 +246,25 @@ def write_message(
     ) -> None:
         """Write a message to the bag file.
 
+        Automatically creates the connection (and schema) if it doesn't exist.
+        If the connection was pre-registered with add_connection(), uses that schema.
+
         Args:
             topic: The topic name.
             timestamp: The timestamp in nanoseconds since epoch.
             message: The message to write.
         """
         message_type = type(message)
 
-        # Ensure connection exists
-        conn_id = self.add_connection(topic, message_type)
+        # Check if connection already exists (may have been pre-registered)
+        if topic in self._topics:
+            conn_id = self._topics[topic]
+        else:
+            # Auto-create connection from message type
+            conn_id = self.add_connection(topic, schema=SchemaText(
+               name=message_type.__msg_name__,
+               text=self._schema_encoder.encode(message_type).decode('utf-8'),
+            ))
 
         # Update chunk time bounds
         if self._chunk_start_time is None:
@@ -234,7 +278,7 @@ def write_message(
         msg_offset = self._chunk_buffer.size()
 
         # Serialize the message
-        serializer = self._get_serializer(message_type)
+        serializer = self._topic_serializers[topic]
         encoder = RosMsgEncoder()
         serializer(encoder, message)
         data = encoder.save()

src/pybag/cli/main.py

@@ -1,6 +1,7 @@
 import argparse
 
 from pybag.cli import (
+    mcap_convert,
     mcap_filter,
     mcap_info,
     mcap_merge,
@@ -23,6 +24,7 @@ def build_parser() -> argparse.ArgumentParser:
     subparsers = parser.add_subparsers(dest="command")
 
     # TODO: Have some of entrypoint registration?
+    mcap_convert.add_parser(subparsers)
     mcap_filter.add_parser(subparsers)
     mcap_merge.add_parser(subparsers)
     mcap_info.add_parser(subparsers)