Add ros1msg schema decoder

src/pybag/schema/__init__.py

@@ -20,6 +20,8 @@
     'uint32': int,
     'int64': int,
     'uint64': int,
+    'time': int,
+    'duration': int,
 }
 STRING_TYPE_MAP = {
     'string': str,

src/pybag/schema/ros1msg.py

@@ -0,0 +1,126 @@
+import logging
+import re
+from typing import Any
+
+from pybag.mcap.records import SchemaRecord
+from pybag.schema import (
+    PRIMITIVE_TYPE_MAP,
+    Array,
+    Complex,
+    Primitive,
+    Schema,
+    SchemaConstant,
+    SchemaDecoder,
+    SchemaEntry,
+    SchemaField,
+    SchemaFieldType,
+    Sequence,
+    String,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Ros1MsgError(Exception):
+    """Exception raised for errors in the ROS1 message parsing."""
+
+
+class Ros1MsgSchemaDecoder(SchemaDecoder):
+    def __init__(self) -> None:
+        self._cache: dict[int, tuple[Schema, dict[str, Schema]]] = {}
+
+    def _remove_inline_comment(self, line: str) -> str:
+        in_single = False
+        in_double = False
+        for i, ch in enumerate(line):
+            if ch == "'" and not in_double:
+                in_single = not in_single
+            elif ch == '"' and not in_single:
+                in_double = not in_double
+            elif ch == '#' and not in_single and not in_double:
+                return line[:i].rstrip()
+        return line.strip()
+
+    def _parse_value(self, field_type: SchemaFieldType, raw_value: str) -> Any:
+        if isinstance(field_type, Primitive):
+            return PRIMITIVE_TYPE_MAP[field_type.type](raw_value)
+        if isinstance(field_type, String):
+            return raw_value.strip('"') if raw_value.startswith('"') else raw_value.strip("'")
+        raise Ros1MsgError('Constants must be primitive or string types')
+
+    def _parse_field_type(self, field_raw_type: str, package_name: str) -> SchemaFieldType:
+        if array_match := re.match(r'(.*)\[(.*)\]$', field_raw_type):
+            element_raw, length_spec = array_match.groups()
+            element_field = self._parse_field_type(element_raw, package_name)
+            if length_spec == '':
+                return Sequence(element_field)
+            length = int(length_spec)
+            return Array(element_field, length)
+
+        if field_raw_type == 'string':
+            return String('string')
+
+        if field_raw_type in PRIMITIVE_TYPE_MAP:
+            return Primitive(field_raw_type)
+
+        if field_raw_type == 'Header':
+            field_raw_type = 'std_msgs/Header'
+        elif '/' not in field_raw_type:
+            field_raw_type = f'{package_name}/{field_raw_type}'
+        return Complex(field_raw_type)
+
+    def _parse_field(self, field: str, package_name: str) -> tuple[str, SchemaEntry]:
+        if '=' in field:
+            if not (match := re.match(r'(\S+)\s+([A-Z][A-Z0-9_]*)=(.+)$', field)):
+                raise Ros1MsgError(f'Invalid constant definition: {field}')
+            field_raw_type, field_raw_name, raw_value = match.groups()
+            schema_type = self._parse_field_type(field_raw_type, package_name)
+            value = self._parse_value(schema_type, raw_value.strip())
+            return field_raw_name, SchemaConstant(schema_type, value)
+
+        if not (match := re.match(r'(\S+)\s+([a-z][a-z0-9_]*)$', field)):
+            raise Ros1MsgError(f'Invalid field definition: {field}')
+        field_raw_type, field_raw_name = match.groups()
+
+        if '__' in field_raw_name:
+            raise Ros1MsgError('Field name cannot contain double underscore "__"')
+        if field_raw_name.endswith('_'):
+            raise Ros1MsgError('Field name cannot end with "_"')
+
+        schema_type = self._parse_field_type(field_raw_type, package_name)
+        return field_raw_name, SchemaField(schema_type)
+
+    def parse_schema(self, schema: SchemaRecord) -> tuple[Schema, dict[str, Schema]]:
+        if schema.id in self._cache:
+            return self._cache[schema.id]
+
+        assert schema.encoding == 'ros1msg'
+        logger.debug(f'Parsing schema: {schema.name}')
+        package_name = schema.name.split('/')[0]
+        msg = schema.data.decode('utf-8')
+
+        lines = [self._remove_inline_comment(line) for line in msg.split('\n')]
+        lines = [line for line in lines if line]
+        msg = '\n'.join(lines)
+
+        msg_parts = [m.strip() for m in msg.split('=' * 80)]
+
+        msg_schema: dict[str, SchemaEntry] = {}
+        main_fields = [m.strip() for m in msg_parts[0].split('\n') if m.strip()]
+        for raw_field in main_fields:
+            field_name, field = self._parse_field(raw_field, package_name)
+            msg_schema[field_name] = field
+
+        sub_msg_schemas: dict[str, Schema] = {}
+        for sub_msg in msg_parts[1:]:
+            sub_msg_name = sub_msg.split('\n')[0].strip()[5:]
+            sub_msg_fields = [m.strip() for m in sub_msg.split('\n')[1:] if m]
+            sub_msg_schema: dict[str, SchemaEntry] = {}
+            for raw_field in sub_msg_fields:
+                field_name, field = self._parse_field(raw_field, package_name)
+                sub_msg_schema[field_name] = field
+            sub_msg_schemas[sub_msg_name] = Schema(sub_msg_name, sub_msg_schema)
+
+        result = Schema(schema.name, msg_schema), sub_msg_schemas
+        self._cache[schema.id] = result
+        return result

tests/schema/test_ros1msg_decoder.py

@@ -0,0 +1,157 @@
+from pybag.schema.ros1msg import (
+    Array,
+    Complex,
+    Primitive,
+    Ros1MsgSchemaDecoder,
+    Schema,
+    SchemaConstant,
+    SchemaField,
+    Sequence,
+)
+from pybag.mcap.records import SchemaRecord
+
+
+def test_parse_primitive_field():
+    schema_text = "int32 my_int\n"
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/Primitive",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    ros1_schema, sub_schemas = Ros1MsgSchemaDecoder().parse_schema(schema)
+
+    assert isinstance(ros1_schema, Schema)
+    assert ros1_schema.name == "pkg/Primitive"
+    assert len(ros1_schema.fields) == 1
+
+    assert "my_int" in ros1_schema.fields
+    field = ros1_schema.fields["my_int"]
+    assert isinstance(field, SchemaField)
+
+    assert isinstance(field.type, Primitive)
+    assert field.type.type == "int32"
+
+    assert sub_schemas == {}
+
+
+def test_parse_unbounded_sequence_field():
+    schema_text = "int32[] values\n"
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/SeqArray",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    ros1_schema, sub_schemas = Ros1MsgSchemaDecoder().parse_schema(schema)
+
+    assert isinstance(ros1_schema, Schema)
+    assert ros1_schema.name == "pkg/SeqArray"
+    assert len(ros1_schema.fields) == 1
+
+    assert "values" in ros1_schema.fields
+    field = ros1_schema.fields["values"]
+    assert isinstance(field, SchemaField)
+
+    assert isinstance(field.type, Sequence)
+
+    assert isinstance(field.type.type, Primitive)
+    assert field.type.type.type == "int32"
+
+    assert sub_schemas == {}
+
+
+def test_parse_static_array_field():
+    schema_text = "int32[5] values\n"
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/StaticArray",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    ros1_schema, sub_schemas = Ros1MsgSchemaDecoder().parse_schema(schema)
+
+    assert isinstance(ros1_schema, Schema)
+    assert ros1_schema.name == "pkg/StaticArray"
+    assert len(ros1_schema.fields) == 1
+
+    assert "values" in ros1_schema.fields
+    field = ros1_schema.fields["values"]
+    assert isinstance(field, SchemaField)
+
+    assert isinstance(field.type, Array)
+    assert field.type.length == 5
+
+    assert isinstance(field.type.type, Primitive)
+    assert field.type.type.type == "int32"
+
+    assert sub_schemas == {}
+
+
+def test_parse_constant():
+    schema_text = "int32 VALUE=5\n"
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/Const",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    ros1_schema, sub_schemas = Ros1MsgSchemaDecoder().parse_schema(schema)
+
+    assert isinstance(ros1_schema, Schema)
+    assert ros1_schema.name == "pkg/Const"
+    assert len(ros1_schema.fields) == 1
+
+    assert "VALUE" in ros1_schema.fields
+    field = ros1_schema.fields["VALUE"]
+    assert isinstance(field, SchemaConstant)
+    assert field.value == 5
+
+    assert isinstance(field.type, Primitive)
+    assert field.type.type == "int32"
+
+    assert sub_schemas == {}
+
+
+def test_parse_complex_field():
+    schema_text = (
+        "geometry_msgs/Point point\n"
+        + "=" * 80
+        + "\nMSG: geometry_msgs/Point\nfloat64 x\nfloat64 y\nfloat64 z\n"
+    )
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/UsePoint",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    ros1_schema, sub_schemas = Ros1MsgSchemaDecoder().parse_schema(schema)
+
+    assert isinstance(ros1_schema, Schema)
+    assert ros1_schema.name == "pkg/UsePoint"
+    assert len(ros1_schema.fields) == 1
+
+    field = ros1_schema.fields["point"]
+    assert isinstance(field, SchemaField)
+    assert isinstance(field.type, Complex)
+    assert field.type.type == "geometry_msgs/Point"
+
+    assert "geometry_msgs/Point" in sub_schemas
+    point_schema = sub_schemas["geometry_msgs/Point"]
+    assert isinstance(point_schema, Schema)
+    assert len(point_schema.fields) == 3
+
+
+def test_parse_caches_schema():
+    schema_text = "int32 my_int\n"
+    schema = SchemaRecord(
+        id=1,
+        name="pkg/Primitive",
+        encoding="ros1msg",
+        data=schema_text.encode("utf-8"),
+    )
+    decoder = Ros1MsgSchemaDecoder()
+    first = decoder.parse_schema(schema)
+    second = decoder.parse_schema(schema)
+
+    assert first is second