Cache schema / channel IDs in Python Writer

python/mcap/mcap/writer.py

@@ -1,4 +1,6 @@
+import functools
 import struct
+import weakref
 import zlib
 from collections import defaultdict
 from enum import Enum, Flag, auto
@@ -60,6 +62,25 @@ class IndexType(Flag):
     ALL = ATTACHMENT | CHUNK | MESSAGE | METADATA
 
 
+def _cache_method(func):
+    """
+    Helper decorator for caching the output of a function (sometimes called memoization) based on
+    the arguments passed to the function, but ignoring self
+
+    :param func: The function to decorate
+    """
+
+    @functools.lru_cache(maxsize=None, typed=False)
+    def _func(_self, *args, **kwargs):
+        return func(_self(), *args, **kwargs)
+
+    @functools.wraps(func)
+    def inner(self, *args, **kwargs):
+        return _func(weakref.ref(self), *args, **kwargs)
+
+    return inner
+
+
 class Writer:
     """
     Writes MCAP data.
@@ -357,6 +378,7 @@ def finish(self):
         if self.__should_close:
             self.__stream.close()
 
+    @_cache_method
     def register_channel(
         self,
         topic: str,
@@ -391,6 +413,7 @@ def register_channel(
             channel.write(self.__record_builder)
         return channel_id
 
+    @_cache_method
     def register_schema(self, name: str, encoding: str, data: bytes):
         """
         Registers a new message schema. Returns the new integer schema id.

python/mcap/tests/test_writer.py

@@ -7,8 +7,7 @@
 
 import lz4.frame
 import pytest
-
-from mcap.records import Chunk, ChunkIndex, Statistics
+from mcap.records import Channel, Chunk, ChunkIndex, Schema, Statistics
 from mcap.stream_reader import StreamReader
 from mcap.writer import CompressionType, Writer
 
@@ -52,6 +51,68 @@ def generate_sample_data(compression: CompressionType):
     yield file
 
 
+@contextlib.contextmanager
+def generate_mcap_schemas_channels(compression: CompressionType):
+    file = TemporaryFile("w+b")
+    writer = Writer(file, compression=compression)
+    writer.start(library="test")
+    schema_id = writer.register_schema(
+        name="sample",
+        encoding="jsonschema",
+        data=json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "sample": {
+                        "type": "string",
+                    }
+                },
+            }
+        ).encode(),
+    )
+    # Schema written twice — this should be cached and there should only be one channel record
+    # written to the final MCAP file
+    schema_id = writer.register_schema(
+        name="sample",
+        encoding="jsonschema",
+        data=json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "sample": {
+                        "type": "string",
+                    }
+                },
+            }
+        ).encode(),
+    )
+
+    channel_id = writer.register_channel(
+        schema_id=schema_id,
+        topic="sample_topic",
+        message_encoding="json",
+    )
+    # Channel written twice — this should be cached and there should only be one channel record
+    # written to the final MCAP file
+    channel_id = writer.register_channel(
+        schema_id=schema_id,
+        topic="sample_topic",
+        message_encoding="json",
+    )
+
+    writer.add_message(
+        channel_id=channel_id,
+        log_time=0,
+        data=json.dumps({"sample": "test"}).encode("utf-8"),
+        publish_time=0,
+    )
+
+    writer.finish()
+    file.seek(0)
+
+    yield file
+
+
 def test_lz4_chunks():
     """tests that compression metadata is correctly written to chunks and chunk indices."""
     chunks: List[Chunk] = []
@@ -115,3 +176,18 @@ def test_out_of_order_messages():
     chunk_index = next(r for r in records if isinstance(r, ChunkIndex))
     assert chunk_index.message_start_time == 0
     assert chunk_index.message_end_time == 100
+
+
+def test_schema_channel_caching():
+    """tests that schema and channel records are cached and not written more than once."""
+    schemas: List[Schema] = []
+    channels: List[Channel] = []
+    with generate_mcap_schemas_channels(CompressionType.LZ4) as t:
+        for record in StreamReader(t, emit_chunks=True).records:
+            if isinstance(record, Schema):
+                schemas.append(record)
+            elif isinstance(record, Channel):
+                channels.append(record)
+
+    assert len(schemas) == 1
+    assert len(channels) == 1