Add rosbag2 mcap storage reader
This commit is contained in:
Marko Durkovic
@@ -5,10 +5,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from rosbags.highlevel import AnyReader, AnyReaderError
|
||||
from rosbags.interfaces import Connection
|
||||
from rosbags.rosbag1 import Writer as Writer1
|
||||
from rosbags.rosbag2 import Writer as Writer2
|
||||
|
||||
@@ -200,3 +202,61 @@ def test_anyreader2(bags2: list[Path]) -> None: # pylint: disable=redefined-out
|
||||
assert nxt[0].topic == '/topic1'
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
|
||||
def test_anyreader2_autoregister(bags2: list[Path]) -> None: # pylint: disable=redefined-outer-name
|
||||
"""Test AnyReader on rosbag2."""
|
||||
|
||||
class MockReader:
|
||||
"""Mock reader."""
|
||||
|
||||
# pylint: disable=too-few-public-methods
|
||||
|
||||
def __init__(self, paths: list[Path]):
|
||||
"""Initialize mock."""
|
||||
_ = paths
|
||||
self.metadata = {'storage_identifier': 'mcap'}
|
||||
self.connections = [
|
||||
Connection(
|
||||
1,
|
||||
'/foo',
|
||||
'test_msg/msg/Foo',
|
||||
'string foo',
|
||||
'msg',
|
||||
0,
|
||||
None, # type: ignore
|
||||
self,
|
||||
),
|
||||
Connection(
|
||||
2,
|
||||
'/bar',
|
||||
'test_msg/msg/Bar',
|
||||
'module test_msgs { module msg { struct Bar {string bar;}; }; };',
|
||||
'idl',
|
||||
0,
|
||||
None, # type: ignore
|
||||
self,
|
||||
),
|
||||
Connection(
|
||||
3,
|
||||
'/baz',
|
||||
'test_msg/msg/Baz',
|
||||
'',
|
||||
'',
|
||||
0,
|
||||
None, # type: ignore
|
||||
self,
|
||||
),
|
||||
]
|
||||
|
||||
def open(self) -> None:
|
||||
"""Unused."""
|
||||
|
||||
with patch('rosbags.highlevel.anyreader.Reader2', MockReader), \
|
||||
patch('rosbags.highlevel.anyreader.register_types') as mock_register_types:
|
||||
AnyReader([bags2[0]]).open()
|
||||
mock_register_types.assert_called_once()
|
||||
assert mock_register_types.call_args[0][0] == {
|
||||
'test_msg/msg/Foo': ([], [('foo', (1, 'string'))]),
|
||||
'test_msgs/msg/Bar': ([], [('bar', (1, 'string'))]),
|
||||
}
|
||||
|
||||
@@ -7,6 +7,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import struct
|
||||
from io import BytesIO
|
||||
from itertools import groupby
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest import mock
|
||||
@@ -19,6 +22,8 @@ from rosbags.rosbag2 import Reader, ReaderError, Writer
|
||||
from .test_serde import MSG_JOINT, MSG_MAGN, MSG_MAGN_BIG, MSG_POLY
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import BinaryIO, Iterable
|
||||
|
||||
from _pytest.fixtures import SubRequest
|
||||
|
||||
METADATA = """
|
||||
@@ -320,3 +325,427 @@ def test_failure_cases(tmp_path: Path) -> None:
|
||||
with pytest.raises(ReaderError, match='not open database'), \
|
||||
Reader(tmp_path) as reader:
|
||||
next(reader.messages())
|
||||
|
||||
|
||||
def write_record(bio: BinaryIO, opcode: int, records: Iterable[bytes]) -> None:
|
||||
"""Write record."""
|
||||
data = b''.join(records)
|
||||
bio.write(bytes([opcode]) + struct.pack('<Q', len(data)) + data)
|
||||
|
||||
|
||||
def make_string(text: str) -> bytes:
|
||||
"""Serialize string."""
|
||||
data = text.encode()
|
||||
return struct.pack('<I', len(data)) + data
|
||||
|
||||
|
||||
MCAP_HEADER = b'\x89MCAP0\r\n'
|
||||
|
||||
SCHEMAS = [
|
||||
(
|
||||
0x03,
|
||||
(
|
||||
struct.pack('<H', 1),
|
||||
make_string('geometry_msgs/msg/Polygon'),
|
||||
make_string('ros2msg'),
|
||||
make_string('string foo'),
|
||||
),
|
||||
),
|
||||
(
|
||||
0x03,
|
||||
(
|
||||
struct.pack('<H', 2),
|
||||
make_string('sensor_msgs/msg/MagneticField'),
|
||||
make_string('ros2msg'),
|
||||
make_string('string foo'),
|
||||
),
|
||||
),
|
||||
(
|
||||
0x03,
|
||||
(
|
||||
struct.pack('<H', 3),
|
||||
make_string('trajectory_msgs/msg/JointTrajectory'),
|
||||
make_string('ros2msg'),
|
||||
make_string('string foo'),
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
CHANNELS = [
|
||||
(
|
||||
0x04,
|
||||
(
|
||||
struct.pack('<H', 1),
|
||||
struct.pack('<H', 1),
|
||||
make_string('/poly'),
|
||||
make_string('cdr'),
|
||||
make_string(''),
|
||||
),
|
||||
),
|
||||
(
|
||||
0x04,
|
||||
(
|
||||
struct.pack('<H', 2),
|
||||
struct.pack('<H', 2),
|
||||
make_string('/magn'),
|
||||
make_string('cdr'),
|
||||
make_string(''),
|
||||
),
|
||||
),
|
||||
(
|
||||
0x04,
|
||||
(
|
||||
struct.pack('<H', 3),
|
||||
struct.pack('<H', 3),
|
||||
make_string('/joint'),
|
||||
make_string('cdr'),
|
||||
make_string(''),
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
params=['unindexed', 'partially_indexed', 'indexed', 'chunked_unindexed', 'chunked_indexed'],
|
||||
)
|
||||
def bag_mcap(request: SubRequest, tmp_path: Path) -> Path:
|
||||
"""Manually contruct mcap bag."""
|
||||
# pylint: disable=too-many-locals
|
||||
# pylint: disable=too-many-statements
|
||||
(tmp_path / 'metadata.yaml').write_text(
|
||||
METADATA.format(
|
||||
extension='.mcap',
|
||||
compression_format='""',
|
||||
compression_mode='""',
|
||||
).replace('sqlite3', 'mcap'),
|
||||
)
|
||||
|
||||
path = tmp_path / 'db.db3.mcap'
|
||||
bio: BinaryIO
|
||||
messages: list[tuple[int, int, int]] = []
|
||||
chunks = []
|
||||
with path.open('wb') as bio:
|
||||
realbio = bio
|
||||
bio.write(MCAP_HEADER)
|
||||
write_record(bio, 0x01, (make_string('ros2'), make_string('test_mcap')))
|
||||
|
||||
if request.param.startswith('chunked'):
|
||||
bio = BytesIO()
|
||||
messages = []
|
||||
|
||||
write_record(bio, *SCHEMAS[0])
|
||||
write_record(bio, *CHANNELS[0])
|
||||
messages.append((1, 666, bio.tell()))
|
||||
write_record(
|
||||
bio,
|
||||
0x05,
|
||||
(
|
||||
struct.pack('<H', 1),
|
||||
struct.pack('<I', 1),
|
||||
struct.pack('<Q', 666),
|
||||
struct.pack('<Q', 666),
|
||||
MSG_POLY[0],
|
||||
),
|
||||
)
|
||||
|
||||
if request.param.startswith('chunked'):
|
||||
assert isinstance(bio, BytesIO)
|
||||
chunk_start = realbio.tell()
|
||||
compression = make_string('')
|
||||
uncompressed_size = struct.pack('<Q', len(bio.getbuffer()))
|
||||
compressed_size = struct.pack('<Q', len(bio.getbuffer()))
|
||||
write_record(
|
||||
realbio,
|
||||
0x06,
|
||||
(
|
||||
struct.pack('<Q', 666),
|
||||
struct.pack('<Q', 666),
|
||||
uncompressed_size,
|
||||
struct.pack('<I', 0),
|
||||
compression,
|
||||
compressed_size,
|
||||
bio.getbuffer(),
|
||||
),
|
||||
)
|
||||
message_index_offsets = []
|
||||
message_index_start = realbio.tell()
|
||||
for channel_id, group in groupby(messages, key=lambda x: x[0]):
|
||||
message_index_offsets.append((channel_id, realbio.tell()))
|
||||
tpls = [y for x in group for y in x[1:]]
|
||||
write_record(
|
||||
realbio,
|
||||
0x07,
|
||||
(
|
||||
struct.pack('<H', channel_id),
|
||||
struct.pack('<I', 8 * len(tpls)),
|
||||
struct.pack('<' + 'Q' * len(tpls), *tpls),
|
||||
),
|
||||
)
|
||||
chunk = [
|
||||
struct.pack('<Q', 666),
|
||||
struct.pack('<Q', 666),
|
||||
struct.pack('<Q', chunk_start),
|
||||
struct.pack('<Q', message_index_start - chunk_start),
|
||||
struct.pack('<I', 10 * len(message_index_offsets)),
|
||||
*(struct.pack('<HQ', *x) for x in message_index_offsets),
|
||||
struct.pack('<Q',
|
||||
realbio.tell() - message_index_start),
|
||||
compression,
|
||||
compressed_size,
|
||||
uncompressed_size,
|
||||
]
|
||||
chunks.append(chunk)
|
||||
bio = BytesIO()
|
||||
messages = []
|
||||
|
||||
write_record(bio, *SCHEMAS[1])
|
||||
write_record(bio, *CHANNELS[1])
|
||||
messages.append((2, 708, bio.tell()))
|
||||
write_record(
|
||||
bio,
|
||||
0x05,
|
||||
(
|
||||
struct.pack('<H', 2),
|
||||
struct.pack('<I', 1),
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', 708),
|
||||
MSG_MAGN[0],
|
||||
),
|
||||
)
|
||||
messages.append((2, 708, bio.tell()))
|
||||
write_record(
|
||||
bio,
|
||||
0x05,
|
||||
(
|
||||
struct.pack('<H', 2),
|
||||
struct.pack('<I', 2),
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', 708),
|
||||
MSG_MAGN_BIG[0],
|
||||
),
|
||||
)
|
||||
|
||||
write_record(bio, *SCHEMAS[2])
|
||||
write_record(bio, *CHANNELS[2])
|
||||
messages.append((3, 708, bio.tell()))
|
||||
write_record(
|
||||
bio,
|
||||
0x05,
|
||||
(
|
||||
struct.pack('<H', 3),
|
||||
struct.pack('<I', 1),
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', 708),
|
||||
MSG_JOINT[0],
|
||||
),
|
||||
)
|
||||
|
||||
if request.param.startswith('chunked'):
|
||||
assert isinstance(bio, BytesIO)
|
||||
chunk_start = realbio.tell()
|
||||
compression = make_string('')
|
||||
uncompressed_size = struct.pack('<Q', len(bio.getbuffer()))
|
||||
compressed_size = struct.pack('<Q', len(bio.getbuffer()))
|
||||
write_record(
|
||||
realbio,
|
||||
0x06,
|
||||
(
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', 708),
|
||||
uncompressed_size,
|
||||
struct.pack('<I', 0),
|
||||
compression,
|
||||
compressed_size,
|
||||
bio.getbuffer(),
|
||||
),
|
||||
)
|
||||
message_index_offsets = []
|
||||
message_index_start = realbio.tell()
|
||||
for channel_id, group in groupby(messages, key=lambda x: x[0]):
|
||||
message_index_offsets.append((channel_id, realbio.tell()))
|
||||
tpls = [y for x in group for y in x[1:]]
|
||||
write_record(
|
||||
realbio,
|
||||
0x07,
|
||||
(
|
||||
struct.pack('<H', channel_id),
|
||||
struct.pack('<I', 8 * len(tpls)),
|
||||
struct.pack('<' + 'Q' * len(tpls), *tpls),
|
||||
),
|
||||
)
|
||||
chunk = [
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<Q', chunk_start),
|
||||
struct.pack('<Q', message_index_start - chunk_start),
|
||||
struct.pack('<I', 10 * len(message_index_offsets)),
|
||||
*(struct.pack('<HQ', *x) for x in message_index_offsets),
|
||||
struct.pack('<Q',
|
||||
realbio.tell() - message_index_start),
|
||||
compression,
|
||||
compressed_size,
|
||||
uncompressed_size,
|
||||
]
|
||||
chunks.append(chunk)
|
||||
bio = realbio
|
||||
messages = []
|
||||
|
||||
if request.param in ['indexed', 'partially_indexed', 'chunked_indexed']:
|
||||
summary_start = bio.tell()
|
||||
for schema in SCHEMAS:
|
||||
write_record(bio, *schema)
|
||||
if request.param != 'partially_indexed':
|
||||
for channel in CHANNELS:
|
||||
write_record(bio, *channel)
|
||||
if request.param == 'chunked_indexed':
|
||||
for chunk in chunks:
|
||||
write_record(bio, 0x08, chunk)
|
||||
|
||||
summary_offset_start = 0
|
||||
write_record(bio, 0x0a, (b'ignored',))
|
||||
write_record(
|
||||
bio,
|
||||
0x0b,
|
||||
(
|
||||
struct.pack('<Q', 4),
|
||||
struct.pack('<H', 3),
|
||||
struct.pack('<I', 3),
|
||||
struct.pack('<I', 0),
|
||||
struct.pack('<I', 0),
|
||||
struct.pack('<I', 0 if request.param == 'indexed' else 1),
|
||||
struct.pack('<Q', 666),
|
||||
struct.pack('<Q', 708),
|
||||
struct.pack('<I', 0),
|
||||
),
|
||||
)
|
||||
write_record(bio, 0x0d, (b'ignored',))
|
||||
write_record(bio, 0xff, (b'ignored',))
|
||||
else:
|
||||
summary_start = 0
|
||||
summary_offset_start = 0
|
||||
|
||||
write_record(
|
||||
bio,
|
||||
0x02,
|
||||
(
|
||||
struct.pack('<Q', summary_start),
|
||||
struct.pack('<Q', summary_offset_start),
|
||||
struct.pack('<I', 0),
|
||||
),
|
||||
)
|
||||
bio.write(MCAP_HEADER)
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
def test_reader_mcap(bag_mcap: Path) -> None:
|
||||
"""Test reader and deserializer on simple bag."""
|
||||
with Reader(bag_mcap) as reader:
|
||||
assert reader.duration == 43
|
||||
assert reader.start_time == 666
|
||||
assert reader.end_time == 709
|
||||
assert reader.message_count == 4
|
||||
if reader.compression_mode:
|
||||
assert reader.compression_format == 'zstd'
|
||||
assert [x.id for x in reader.connections] == [1, 2, 3]
|
||||
assert [*reader.topics.keys()] == ['/poly', '/magn', '/joint']
|
||||
gen = reader.messages()
|
||||
|
||||
connection, timestamp, rawdata = next(gen)
|
||||
assert connection.topic == '/poly'
|
||||
assert connection.msgtype == 'geometry_msgs/msg/Polygon'
|
||||
assert timestamp == 666
|
||||
assert rawdata == MSG_POLY[0]
|
||||
|
||||
for idx in range(2):
|
||||
connection, timestamp, rawdata = next(gen)
|
||||
assert connection.topic == '/magn'
|
||||
assert connection.msgtype == 'sensor_msgs/msg/MagneticField'
|
||||
assert timestamp == 708
|
||||
assert rawdata == [MSG_MAGN, MSG_MAGN_BIG][idx][0]
|
||||
|
||||
connection, timestamp, rawdata = next(gen)
|
||||
assert connection.topic == '/joint'
|
||||
assert connection.msgtype == 'trajectory_msgs/msg/JointTrajectory'
|
||||
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
|
||||
def test_message_filters_mcap(bag_mcap: Path) -> None:
|
||||
"""Test reader filters messages."""
|
||||
with Reader(bag_mcap) as reader:
|
||||
magn_connections = [x for x in reader.connections if x.topic == '/magn']
|
||||
gen = reader.messages(connections=magn_connections)
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/magn'
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/magn'
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
gen = reader.messages(start=667)
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/magn'
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/magn'
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/joint'
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
gen = reader.messages(stop=667)
|
||||
connection, _, _ = next(gen)
|
||||
assert connection.topic == '/poly'
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
gen = reader.messages(connections=magn_connections, stop=667)
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
gen = reader.messages(start=666, stop=666)
|
||||
with pytest.raises(StopIteration):
|
||||
next(gen)
|
||||
|
||||
|
||||
def test_bag_mcap_files(tmp_path: Path) -> None:
|
||||
"""Test bad mcap files."""
|
||||
(tmp_path / 'metadata.yaml').write_text(
|
||||
METADATA.format(
|
||||
extension='.mcap',
|
||||
compression_format='""',
|
||||
compression_mode='""',
|
||||
).replace('sqlite3', 'mcap'),
|
||||
)
|
||||
|
||||
path = tmp_path / 'db.db3.mcap'
|
||||
path.touch()
|
||||
reader = Reader(tmp_path)
|
||||
path.unlink()
|
||||
with pytest.raises(ReaderError, match='Could not open'):
|
||||
reader.open()
|
||||
|
||||
path.touch()
|
||||
with pytest.raises(ReaderError, match='seems to be empty'):
|
||||
Reader(tmp_path).open()
|
||||
|
||||
path.write_bytes(b'xxxxxxxx')
|
||||
with pytest.raises(ReaderError, match='magic is invalid'):
|
||||
Reader(tmp_path).open()
|
||||
|
||||
path.write_bytes(b'\x89MCAP0\r\n\xFF')
|
||||
with pytest.raises(ReaderError, match='Unexpected record'):
|
||||
Reader(tmp_path).open()
|
||||
|
||||
with path.open('wb') as bio:
|
||||
bio.write(b'\x89MCAP0\r\n')
|
||||
write_record(bio, 0x01, (make_string('ros1'), make_string('test_mcap')))
|
||||
with pytest.raises(ReaderError, match='Profile is not'):
|
||||
Reader(tmp_path).open()
|
||||
|
||||
with path.open('wb') as bio:
|
||||
bio.write(b'\x89MCAP0\r\n')
|
||||
write_record(bio, 0x01, (make_string('ros2'), make_string('test_mcap')))
|
||||
with pytest.raises(ReaderError, match='File end magic is invalid'):
|
||||
Reader(tmp_path).open()
|
||||
|
||||
Reference in New Issue
Block a user