apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add rosbag2 support

Browse Source
This commit is contained in:
Marko Durkovic 2021-05-02 14:49:33 +02:00
parent abd0c1fa73
commit ffacb7602c
11 changed files with 997 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
README.rst
View File

@ -32,6 +32,20 @@ Rosbags is published on PyPI and does not have any special dependencies. Simply
pip install rosbags
Read and deserialize rosbag2 messages:
.. code-block:: python
from rosbags.rosbag2 import Reader
from rosbags.serde import deserialize_cdr
# create reader instance and open for reading
with Reader('/home/ros/rosbag_2020_03_24') as reader:
for topic, msgtype, timestamp, rawdata in reader.messages(['/imu_raw/Imu']):
msg = deserialize_cdr(rawdata, msgtype)
print(msg.header.frame_id)
Documentation
=============

6
docs/api/rosbags.rosbag2.rst Normal file
View File

@ -0,0 +1,6 @@
rosbags.rosbag2
===============
.. automodule:: rosbags.rosbag2
:members:
:show-inheritance:

1
docs/api/rosbags.rst
View File

@ -4,5 +4,6 @@ Rosbags namespace
.. toctree::
:maxdepth: 4
rosbags.rosbag2
rosbags.serde
rosbags.typesys

1
docs/index.rst
View File

@ -12,6 +12,7 @@
topics/typesys
topics/serde
topics/rosbag2
.. toctree::

65
docs/topics/rosbag2.rst Normal file
View File

@ -0,0 +1,65 @@
Rosbag2
=======
The :py:mod:`rosbags.rosbag2` package provides a conformant implementation of rosbag2. It provides read-write access to raw message data saved inside rosbag2 containers, and supports all features present in the C++ reference implementation.
Supported Versions
------------------
All versions up to the current (ROS2 Foxy) version 4 are supported.
Supported Features
------------------
Rosbag2 is a flexible format that supports plugging different serialization methods, compression formats, and storage containers together. The rosbag2 C++ reference implementation is build around plugins that provide serialization, compression, and storage. This project implements all rosbag2 core plugins that are distributed with the C++ reference implementation.
:Serializers:
- cdr (without wstring)
:Compressors:
- zstd
:Storages:
- sqlite3
Writing rosbag2
---------------
Instances of the :py:class:`Writer <rosbags.rosbag2.Writer>` class can create and write to new rosbag2 files. It is usually used as a context manager. Before the first message of a topic can be written, its topic must first be added to the bag. The following example shows the typical usage pattern:
.. code-block:: python
from rosbags.rosbag2 import Writer
from rosbags.serde import serialize_cdr
# create writer instance and open for writing
with Writer('/home/ros/rosbag_2020_03_24') as writer:
# add new topic
topic = '/imu_raw/Imu'
msgtype = 'sensor_msgs/msg/Imu'
writer.add_topic(topic, msgtype, 'cdr')
# serialize and write message
writer.write(topic, timestamp, serialize_cdr(message, msgtype))
Reading rosbag2
---------------
Instances of the :py:class:`Reader <rosbags.rosbag2.Reader>` class are used to read rosbag2 metadata and its contents. Most of the metadata is available on Reader instances right away, messages can only be accessed after the bag has been opened. To this end it is recommended to use the Reader as a context manager. The following example shows the typical usage pattern:
.. code-block:: python
from rosbags.rosbag2 import Reader
from rosbags.serde import deserialize_cdr
# create reader instance and open for reading
with Reader('/home/ros/rosbag_2020_03_24') as reader:
# topic and msgtype information is available on .topics dict
for topic, msgtype in reader.topics.items():
print(topic, msgtype)
# iterate over messages
for topic, msgtype, timestamp, rawdata in reader.messages():
if topic == '/imu_raw/Imu':
msg = deserialize_cdr(rawdata, msgtype)
print(msg.header.frame_id)
# messages() accepts topic filters
for topic, msgtype, rawdata, timestamp in reader.messages(['/imu_raw/Imu']):
msg = deserialize_cdr(rawdata, msgtype)
print(msg.header.frame_id)

18
src/rosbags/rosbag2/__init__.py Normal file
View File

@ -0,0 +1,18 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Rosbags support for rosbag2 files.
Readers and writers provide access to metadata and raw message content saved
in the rosbag2 format.
"""
from .reader import Reader, ReaderError
from .writer import Writer, WriterError
__all__ = [
'Reader',
'ReaderError',
'Writer',
'WriterError',
]

231
src/rosbags/rosbag2/reader.py Normal file
View File

@ -0,0 +1,231 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Rosbag2 reader."""
from __future__ import annotations
import sqlite3
from contextlib import contextmanager
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING
import zstandard
from ruamel.yaml import YAML, YAMLError
if TYPE_CHECKING:
from types import TracebackType
from typing import Any, Generator, Iterable, List, Literal, Optional, Tuple, Type, Union
class ReaderError(Exception):
"""Reader Error."""
@contextmanager
def decompress(path: Path, do_decompress: bool):
"""Transparent rosbag2 database decompression context.
This context manager will yield a path to the decompressed file contents.
Args:
path: Potentially compressed file.
do_decompress: Flag indicating if decompression shall occur.
Yields:
Path of transparently decompressed file.
"""
if do_decompress:
decomp = zstandard.ZstdDecompressor()
with TemporaryDirectory() as tempdir:
dbfile = Path(tempdir, path.stem)
with path.open('rb') as infile, dbfile.open('wb') as outfile:
decomp.copy_stream(infile, outfile)
yield dbfile
else:
yield path
class Reader:
"""Reader for rosbag2 files.
It implements all necessary features to access metadata and message
streams.
Version history:
- Version 1: Initial format.
- Version 2: Changed field sizes in C++ implementation.
- Version 3: Added compression.
- Version 4: Added QoS metadata to topics, changed relative file paths
"""
def __init__(self, path: Union[Path, str]):
"""Open rosbag and check metadata.
Args:
path: Filesystem path to bag.
Raises:
ReaderError: Bag not readable or bag metadata.
"""
path = Path(path)
self.path = Path
self.bio = False
try:
yaml = YAML(typ='safe')
yamlpath = path / 'metadata.yaml'
dct = yaml.load(yamlpath.read_text())
except OSError as err:
raise ReaderError(f'Could not read metadata at {yamlpath}: {err}.') from None
except YAMLError as exc:
raise ReaderError(f'Could not load YAML from {yamlpath}: {exc}') from None
try:
self.metadata = dct['rosbag2_bagfile_information']
if (ver := self.metadata['version']) > 4:
raise ReaderError(f'Rosbag2 version {ver} not supported; please report issue.')
if storageid := self.metadata['storage_identifier'] != 'sqlite3':
raise ReaderError(
f'Storage plugin {storageid!r} not supported; please report issue.',
)
self.paths = [path / Path(x).name for x in self.metadata['relative_file_paths']]
missing = [x for x in self.paths if not x.exists()]
if missing:
raise ReaderError(f'Some database files are missing: {[str(x) for x in missing]!r}')
topics = [x['topic_metadata'] for x in self.metadata['topics_with_message_count']]
noncdr = {y for x in topics if (y := x['serialization_format']) != 'cdr'}
if noncdr:
raise ReaderError(f'Serialization format {noncdr!r} is not supported.')
self.topics = {x['name']: x['type'] for x in topics}
if self.compression_mode and (cfmt := self.compression_format) != 'zstd':
raise ReaderError(f'Compression format {cfmt!r} is not supported.')
except KeyError as exc:
raise ReaderError(f'A metadata key is missing {exc!r}.') from None
def open(self):
"""Open rosbag2."""
# Future storage formats will require file handles.
self.bio = True
def close(self):
"""Close rosbag2."""
# Future storage formats will require file handles.
assert self.bio
self.bio = False
@property
def duration(self) -> int:
"""Duration in nanoseconds between earliest and latest messages."""
return self.metadata['duration']['nanoseconds']
@property
def start_time(self) -> int:
"""Timestamp in nanoseconds of the earliest message."""
return self.metadata['starting_time']['nanoseconds_since_epoch']
@property
def end_time(self) -> int:
"""Timestamp in nanoseconds of the latest message."""
return self.start_time + self.duration
@property
def message_count(self) -> int:
"""Total message count."""
return self.metadata['message_count']
@property
def compression_format(self) -> Optional[str]:
"""Compression format."""
return self.metadata.get('compression_format', None) or None
@property
def compression_mode(self) -> Optional[str]:
"""Compression mode."""
mode = self.metadata.get('compression_mode', '').lower()
return mode if mode != 'none' else None
def messages( # pylint: disable=too-many-locals
self,
topics: Iterable[str] = (),
start: Optional[int] = None,
stop: Optional[int] = None,
) -> Generator[Tuple[str, str, int, bytes], None, None]:
"""Read messages from bag.
Args:
topics: Iterable with topic names to filter for. An empty iterable
yields all messages.
start: Yield only messages at or after this timestamp (ns).
stop: Yield only messages before this timestamp (ns).
Yields:
Tuples of topic name, type, timestamp (ns), and rawdata.
Raises:
ReaderError: Bag not open.
"""
if not self.bio:
raise ReaderError('Rosbag is not open.')
topics = tuple(topics)
for filepath in self.paths:
with decompress(filepath, self.compression_mode == 'file') as path:
conn = sqlite3.connect(f'file:{path}?immutable=1', uri=True)
conn.row_factory = lambda _, x: x
cur = conn.cursor()
cur.execute(
'SELECT count(*) FROM sqlite_master '
'WHERE type="table" AND name IN ("messages", "topics")',
)
if cur.fetchone()[0] != 2:
raise ReaderError(f'Cannot open database {path} or database missing tables.')
query = [
'SELECT topics.name,topics.type,messages.timestamp,messages.data',
'FROM messages JOIN topics ON messages.topic_id=topics.id',
]
args: List[Any] = []
if topics:
query.append(f'WHERE topics.name IN ({",".join("?" for _ in topics)})')
args += topics
if start is not None:
query.append(f'{"AND" if args else "WHERE"} messages.timestamp >= ?')
args.append(start)
if stop is not None:
query.append(f'{"AND" if args else "WHERE"} messages.timestamp < ?')
args.append(stop)
query.append('ORDER BY timestamp')
cur.execute(' '.join(query), args)
if self.compression_mode == 'message':
decomp = zstandard.ZstdDecompressor().decompress
for row in cur:
topic, msgtype, timestamp, data = row
yield topic, msgtype, timestamp, decomp(data)
else:
yield from cur
def __enter__(self) -> Reader:
"""Open rosbag2 when entering contextmanager."""
self.open()
return self
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> Literal[False]:
"""Close rosbag2 when exiting contextmanager."""
self.close()
return False

247
src/rosbags/rosbag2/writer.py Normal file
View File

@ -0,0 +1,247 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Rosbag2 reader."""
from __future__ import annotations
import sqlite3
from enum import IntEnum, auto
from pathlib import Path
from typing import TYPE_CHECKING
import zstandard
from ruamel.yaml import YAML
if TYPE_CHECKING:
from types import TracebackType
from typing import Any, Dict, Literal, Optional, Type, Union
class WriterError(Exception):
"""Writer Error."""
class Writer: # pylint: disable=too-many-instance-attributes
"""Rosbag2 writer.
This class implements writing of rosbag2 files in version 4. It should be
used as a contextmanager.
"""
SQLITE_SCHEMA = """
CREATE TABLE topics(
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
type TEXT NOT NULL,
serialization_format TEXT NOT NULL,
offered_qos_profiles TEXT NOT NULL
);
CREATE TABLE messages(
id INTEGER PRIMARY KEY,
topic_id INTEGER NOT NULL,
timestamp INTEGER NOT NULL,
data BLOB NOT NULL
);
CREATE INDEX timestamp_idx ON messages (timestamp ASC);
"""
class CompressionMode(IntEnum):
"""Compession modes."""
NONE = auto()
FILE = auto()
MESSAGE = auto()
class CompressionFormat(IntEnum):
"""Compession formats."""
ZSTD = auto()
def __init__(self, path: Union[Path, str]):
"""Initialize writer.
Args:
path: Filesystem path to bag.
Raises:
WriterError: Target path exisits already, Writer can only create new rosbags.
"""
path = Path(path)
self.path = path
if path.exists():
raise WriterError(f'{path} exists already, not overwriting.')
self.metapath = path / 'metadata.yaml'
self.dbpath = path / f'{path.name}.db3'
self.compression_mode = ''
self.compression_format = ''
self.compressor: Optional[zstandard.ZstdCompressor] = None
self.topics: Dict[str, Any] = {}
self.conn = None
self.cursor: Optional[sqlite3.Cursor] = None
self.topics = {}
def set_compression(self, mode: CompressionMode, fmt: CompressionFormat):
"""Enable compression on bag.
This function has to be called before opening.
Args:
mode: Compression mode to use, either 'file' or 'message'
fmt: Compressor to use, currently only 'zstd'
Raises:
WriterError: Bag already open.
"""
if self.conn:
raise WriterError(f'Cannot set compression, bag {self.path} already open.')
if mode == self.CompressionMode.NONE:
return
self.compression_mode = mode.name.lower()
self.compression_format = fmt.name.lower()
self.compressor = zstandard.ZstdCompressor()
def open(self):
"""Open rosbag2 for writing.
Create base directory and open database connection.
"""
try:
self.path.mkdir(mode=0o755, parents=True)
except FileExistsError:
raise WriterError(f'{self.path} exists already, not overwriting.') from None
self.conn = sqlite3.connect(f'file:{self.dbpath}', uri=True)
self.conn.executescript(self.SQLITE_SCHEMA)
self.cursor = self.conn.cursor()
def add_topic(
self,
name: str,
typ: str,
serialization_format: str = 'cdr',
offered_qos_profiles: str = '',
):
"""Add a topic.
This function can only be called after opening a bag.
Args:
name: Topic name.
typ: Message type.
serialization_format: Serialization format.
offered_qos_profiles: QOS Profile.
Raises:
WriterError: Bag not open or topic previously registered.
"""
if not self.cursor:
raise WriterError('Bag was not opened.')
if name in self.topics:
raise WriterError(f'Topics can only be added once: {name!r}.')
meta = (len(self.topics) + 1, name, typ, serialization_format, offered_qos_profiles)
self.cursor.execute('INSERT INTO topics VALUES(?, ?, ?, ?, ?)', meta)
self.topics[name] = [*meta, 0]
def write(self, topic: str, timestamp: int, data: bytes):
"""Write message to rosbag2.
Args:
topic: Topic message belongs to.
timestamp: Message timestamp (ns).
data: Serialized message data.
Raises:
WriterError: Bag not open or topic not registered.
"""
if not self.cursor:
raise WriterError('Bag was not opened.')
if topic not in self.topics:
raise WriterError(f'Tried to write to unknown topic {topic!r}.')
if self.compression_mode == 'message':
assert self.compressor
data = self.compressor.compress(data)
tmeta = self.topics[topic]
self.cursor.execute(
'INSERT INTO messages (topic_id, timestamp, data) VALUES(?, ?, ?)',
(tmeta[0], timestamp, data),
)
tmeta[-1] += 1
def close(self):
"""Close rosbag2 after writing.
Closes open database transactions and writes metadata.yaml.
"""
self.cursor.close()
self.cursor = None
duration, start, count = self.conn.execute(
'SELECT max(timestamp) - min(timestamp) + 1, min(timestamp), count(*) FROM messages',
).fetchone()
self.conn.commit()
self.conn.execute('PRAGMA optimize')
self.conn.close()
if self.compression_mode == 'file':
src = self.dbpath
self.dbpath = src.with_suffix(f'.db3.{self.compression_format}')
with src.open('rb') as infile, self.dbpath.open('wb') as outfile:
self.compressor.copy_stream(infile, outfile)
src.unlink()
metadata = {
'rosbag2_bagfile_information': {
'version': 4,
'storage_identifier': 'sqlite3',
'relative_file_paths': [self.dbpath.name],
'duration': {
'nanoseconds': duration,
},
'starting_time': {
'nanoseconds_since_epoch': start,
},
'message_count': count,
'topics_with_message_count': [
{
'topic_metadata': {
'name': x[1],
'type': x[2],
'serialization_format': x[3],
'offered_qos_profiles': x[4],
},
'message_count': x[5],
} for x in self.topics.values()
],
'compression_format': self.compression_format,
'compression_mode': self.compression_mode,
},
}
with self.metapath.open('w') as metafile:
yaml = YAML(typ='safe')
yaml.default_flow_style = False
yaml.dump(metadata, metafile)
def __enter__(self) -> Writer:
"""Open rosbag2 when entering contextmanager."""
self.open()
return self
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> Literal[False]:
"""Close rosbag2 when exiting contextmanager."""
self.close()
return False

278
tests/test_reader.py Normal file
View File

@ -0,0 +1,278 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Reader tests."""
# pylint: disable=redefined-outer-name
from __future__ import annotations
import sqlite3
from pathlib import Path
from typing import TYPE_CHECKING
from unittest import mock
import pytest
import zstandard
from rosbags.rosbag2 import Reader, ReaderError, Writer
from .test_serde import MSG_JOINT, MSG_MAGN, MSG_MAGN_BIG, MSG_POLY
if TYPE_CHECKING:
from _pytest.fixtures import SubRequest
METADATA = """
rosbag2_bagfile_information:
version: 4
storage_identifier: sqlite3
relative_file_paths:
- db.db3{extension}
duration:
nanoseconds: 42
starting_time:
nanoseconds_since_epoch: 666
message_count: 4
topics_with_message_count:
- topic_metadata:
name: /poly
type: geometry_msgs/msg/Polygon
serialization_format: cdr
offered_qos_profiles: ""
message_count: 1
- topic_metadata:
name: /magn
type: sensor_msgs/msg/MagneticField
serialization_format: cdr
offered_qos_profiles: ""
message_count: 2
- topic_metadata:
name: /joint
type: trajectory_msgs/msg/JointTrajectory
serialization_format: cdr
offered_qos_profiles: ""
message_count: 1
compression_format: {compression_format}
compression_mode: {compression_mode}
"""
@pytest.fixture(params=['none', 'file', 'message'])
def bag(request: SubRequest, tmp_path: Path) -> Path:
"""Manually contruct bag."""
(tmp_path / 'metadata.yaml').write_text(
METADATA.format(
extension='' if request.param != 'file' else '.zstd',
compression_format='""' if request.param == 'none' else 'zstd',
compression_mode='""' if request.param == 'none' else request.param.upper(),
),
)
comp = zstandard.ZstdCompressor()
dbpath = tmp_path / 'db.db3'
dbh = sqlite3.connect(dbpath)
dbh.executescript(Writer.SQLITE_SCHEMA)
cur = dbh.cursor()
cur.execute(
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
(1, '/poly', 'geometry_msgs/msg/Polygon', 'cdr', ''),
)
cur.execute(
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
(2, '/magn', 'sensor_msgs/msg/MagneticField', 'cdr', ''),
)
cur.execute(
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
(3, '/joint', 'trajectory_msgs/msg/JointTrajectory', 'cdr', ''),
)
cur.execute(
'INSERT INTO messages VALUES(?, ?, ?, ?)',
(1, 1, 666, MSG_POLY[0] if request.param != 'message' else comp.compress(MSG_POLY[0])),
)
cur.execute(
'INSERT INTO messages VALUES(?, ?, ?, ?)',
(2, 2, 708, MSG_MAGN[0] if request.param != 'message' else comp.compress(MSG_MAGN[0])),
)
cur.execute(
'INSERT INTO messages VALUES(?, ?, ?, ?)',
(
3,
2,
708,
MSG_MAGN_BIG[0] if request.param != 'message' else comp.compress(MSG_MAGN_BIG[0]),
),
)
cur.execute(
'INSERT INTO messages VALUES(?, ?, ?, ?)',
(4, 3, 708, MSG_JOINT[0] if request.param != 'message' else comp.compress(MSG_JOINT[0])),
)
dbh.commit()
if request.param == 'file':
with dbpath.open('rb') as ifh, (tmp_path / 'db.db3.zstd').open('wb') as ofh:
comp.copy_stream(ifh, ofh)
dbpath.unlink()
return tmp_path
def test_reader(bag: Path):
"""Test reader and deserializer on simple bag."""
with Reader(bag) as reader:
assert reader.duration == 42
assert reader.start_time == 666
assert reader.end_time == 708
assert reader.message_count == 4
if reader.compression_mode:
assert reader.compression_format == 'zstd'
gen = reader.messages()
topic, msgtype, timestamp, rawdata = next(gen)
assert topic == '/poly'
assert msgtype == 'geometry_msgs/msg/Polygon'
assert timestamp == 666
assert rawdata == MSG_POLY[0]
for idx in range(2):
topic, msgtype, timestamp, rawdata = next(gen)
assert topic == '/magn'
assert msgtype == 'sensor_msgs/msg/MagneticField'
assert timestamp == 708
assert rawdata == [MSG_MAGN, MSG_MAGN_BIG][idx][0]
topic, msgtype, timestamp, rawdata = next(gen)
assert topic == '/joint'
assert msgtype == 'trajectory_msgs/msg/JointTrajectory'
with pytest.raises(StopIteration):
next(gen)
def test_message_filters(bag: Path):
"""Test reader filters messages."""
with Reader(bag) as reader:
gen = reader.messages(['/magn'])
topic, _, _, _ = next(gen)
assert topic == '/magn'
topic, _, _, _ = next(gen)
assert topic == '/magn'
with pytest.raises(StopIteration):
next(gen)
gen = reader.messages(start=667)
topic, _, _, _ = next(gen)
assert topic == '/magn'
topic, _, _, _ = next(gen)
assert topic == '/magn'
topic, _, _, _ = next(gen)
assert topic == '/joint'
with pytest.raises(StopIteration):
next(gen)
gen = reader.messages(stop=667)
topic, _, _, _ = next(gen)
assert topic == '/poly'
with pytest.raises(StopIteration):
next(gen)
gen = reader.messages(['/magn'], stop=667)
with pytest.raises(StopIteration):
next(gen)
gen = reader.messages(start=666, stop=666)
with pytest.raises(StopIteration):
next(gen)
def test_user_errors(bag: Path):
"""Test user errors."""
reader = Reader(bag)
with pytest.raises(ReaderError, match='Rosbag is not open'):
next(reader.messages())
def test_failure_cases(tmp_path: Path):
"""Test bags with broken fs layout."""
with pytest.raises(ReaderError, match='not read metadata'):
Reader(tmp_path)
metadata = tmp_path / 'metadata.yaml'
metadata.write_text('')
with pytest.raises(ReaderError, match='not read'), \
mock.patch.object(Path, 'read_text', side_effect=PermissionError):
Reader(tmp_path)
metadata.write_text(' invalid:\nthis is not yaml')
with pytest.raises(ReaderError, match='not load YAML from'):
Reader(tmp_path)
metadata.write_text('foo:')
with pytest.raises(ReaderError, match='key is missing'):
Reader(tmp_path)
metadata.write_text(
METADATA.format(
extension='',
compression_format='""',
compression_mode='""',
).replace('version: 4', 'version: 999'),
)
with pytest.raises(ReaderError, match='version 999'):
Reader(tmp_path)
metadata.write_text(
METADATA.format(
extension='',
compression_format='""',
compression_mode='""',
).replace('sqlite3', 'hdf5'),
)
with pytest.raises(ReaderError, match='Storage plugin'):
Reader(tmp_path)
metadata.write_text(
METADATA.format(
extension='',
compression_format='""',
compression_mode='""',
),
)
with pytest.raises(ReaderError, match='files are missing'):
Reader(tmp_path)
(tmp_path / 'db.db3').write_text('')
metadata.write_text(
METADATA.format(
extension='',
compression_format='""',
compression_mode='""',
).replace('cdr', 'bson'),
)
with pytest.raises(ReaderError, match='Serialization format'):
Reader(tmp_path)
metadata.write_text(
METADATA.format(
extension='',
compression_format='"gz"',
compression_mode='"file"',
),
)
with pytest.raises(ReaderError, match='Compression format'):
Reader(tmp_path)
metadata.write_text(
METADATA.format(
extension='',
compression_format='""',
compression_mode='""',
),
)
with pytest.raises(ReaderError, match='not open database'), \
Reader(tmp_path) as reader:
next(reader.messages())

42
tests/test_roundtrip.py Normal file
View File

@ -0,0 +1,42 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Test full data roundtrip."""
from __future__ import annotations
from typing import TYPE_CHECKING
import pytest
from rosbags.rosbag2 import Reader, Writer
from rosbags.serde import deserialize_cdr, serialize_cdr
if TYPE_CHECKING:
from pathlib import Path
@pytest.mark.parametrize('mode', [*Writer.CompressionMode])
def test_roundtrip(mode: Writer.CompressionMode, tmp_path: Path):
"""Test full data roundtrip."""
class Foo: # pylint: disable=too-few-public-methods
"""Dummy class."""
data = 1.25
path = tmp_path / 'rosbag2'
wbag = Writer(path)
wbag.set_compression(mode, wbag.CompressionFormat.ZSTD)
with wbag:
msgtype = 'std_msgs/msg/Float64'
wbag.add_topic('/test', msgtype)
wbag.write('/test', 42, serialize_cdr(Foo, msgtype))
rbag = Reader(path)
with rbag:
gen = rbag.messages()
_, msgtype, _, raw = next(gen)
msg = deserialize_cdr(raw, msgtype)
assert msg.data == Foo.data
with pytest.raises(StopIteration):
next(gen)

94
tests/test_writer.py Normal file
View File

@ -0,0 +1,94 @@
# Copyright 2020-2021 Ternaris.
# SPDX-License-Identifier: Apache-2.0
"""Writer tests."""
from __future__ import annotations
from typing import TYPE_CHECKING
import pytest
from rosbags.rosbag2 import Writer, WriterError
if TYPE_CHECKING:
from pathlib import Path
def test_writer(tmp_path: Path):
"""Test Writer."""
path = (tmp_path / 'rosbag2')
with Writer(path) as bag:
bag.add_topic('/test', 'std_msgs/msg/Int8')
bag.write('/test', 42, b'\x00')
bag.write('/test', 666, b'\x01' * 4096)
assert (path / 'metadata.yaml').exists()
assert (path / 'rosbag2.db3').exists()
size = (path / 'rosbag2.db3').stat().st_size
path = (tmp_path / 'compress_none')
bag = Writer(path)
bag.set_compression(bag.CompressionMode.NONE, bag.CompressionFormat.ZSTD)
with bag:
bag.add_topic('/test', 'std_msgs/msg/Int8')
bag.write('/test', 42, b'\x00')
bag.write('/test', 666, b'\x01' * 4096)
assert (path / 'metadata.yaml').exists()
assert (path / 'compress_none.db3').exists()
assert size == (path / 'compress_none.db3').stat().st_size
path = (tmp_path / 'compress_file')
bag = Writer(path)
bag.set_compression(bag.CompressionMode.FILE, bag.CompressionFormat.ZSTD)
with bag:
bag.add_topic('/test', 'std_msgs/msg/Int8')
bag.write('/test', 42, b'\x00')
bag.write('/test', 666, b'\x01' * 4096)
assert (path / 'metadata.yaml').exists()
assert not (path / 'compress_file.db3').exists()
assert (path / 'compress_file.db3.zstd').exists()
path = (tmp_path / 'compress_message')
bag = Writer(path)
bag.set_compression(bag.CompressionMode.MESSAGE, bag.CompressionFormat.ZSTD)
with bag:
bag.add_topic('/test', 'std_msgs/msg/Int8')
bag.write('/test', 42, b'\x00')
bag.write('/test', 666, b'\x01' * 4096)
assert (path / 'metadata.yaml').exists()
assert (path / 'compress_message.db3').exists()
assert size > (path / 'compress_message.db3').stat().st_size
def test_failure_cases(tmp_path: Path):
"""Test writer failure cases."""
with pytest.raises(WriterError, match='exists'):
Writer(tmp_path)
bag = Writer(tmp_path / 'race')
(tmp_path / 'race').mkdir()
with pytest.raises(WriterError, match='exists'):
bag.open()
bag = Writer(tmp_path / 'compress_after_open')
bag.open()
with pytest.raises(WriterError, match='already open'):
bag.set_compression(bag.CompressionMode.FILE, bag.CompressionFormat.ZSTD)
bag = Writer(tmp_path / 'topic')
with pytest.raises(WriterError, match='was not opened'):
bag.add_topic('/tf', 'tf_msgs/msg/tf2')
bag = Writer(tmp_path / 'write')
with pytest.raises(WriterError, match='was not opened'):
bag.write('/tf', 0, b'')
bag = Writer(tmp_path / 'topic')
bag.open()
bag.add_topic('/tf', 'tf_msgs/msg/tf2')
with pytest.raises(WriterError, match='only be added once'):
bag.add_topic('/tf', 'tf_msgs/msg/tf2')
bag = Writer(tmp_path / 'notopic')
bag.open()
with pytest.raises(WriterError, match='unknown topic'):
bag.write('/test', 42, b'\x00')