Add topic filters to rosbag conversion

This commit is contained in:
Marko Durkovic 2022-07-27 15:39:47 +02:00
parent 05df51aa96
commit b9fd0b014b
3 changed files with 54 additions and 17 deletions

View File

@ -50,6 +50,13 @@ def main() -> None:
type=pathtype(exists=False),
help='destination path for converted rosbag',
)
parser.add_argument(
'--exclude-topic',
action='append',
default=[],
dest='exclude_topics',
help='exclude topic by name',
)
args = parser.parse_args()
if args.dst is not None and (args.src.suffix == '.bag') == (args.dst.suffix == '.bag'):
@ -57,7 +64,7 @@ def main() -> None:
sys.exit(1)
try:
convert(args.src, args.dst)
convert(**args.__dict__)
except ConverterError as err:
print(f'ERROR: {err}') # noqa: T001
sys.exit(1)

View File

@ -21,7 +21,7 @@ from rosbags.typesys.msg import generate_msgdef
if TYPE_CHECKING:
from pathlib import Path
from typing import Any, Optional
from typing import Any, Optional, Sequence
LATCH = """
- history: 3
@ -99,19 +99,25 @@ def downgrade_connection(rconn: Connection) -> Connection:
)
def convert_1to2(src: Path, dst: Path) -> None:
def convert_1to2(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None:
"""Convert Rosbag1 to Rosbag2.
Args:
src: Rosbag1 path.
dst: Rosbag2 path.
exclude_topics: Topics to skip.
Raises:
ConverterError: If all connections are excluded.
"""
with Reader1(src) as reader, Writer2(dst) as writer:
typs: dict[str, Any] = {}
connmap: dict[int, Connection] = {}
for rconn in reader.connections:
connections = [x for x in reader.connections if x.topic not in exclude_topics]
if not connections:
raise ConverterError('No connections left for conversion.')
for rconn in connections:
candidate = upgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag2)
for conn in writer.connections:
@ -132,22 +138,29 @@ def convert_1to2(src: Path, dst: Path) -> None:
typs.update(get_types_from_msg(rconn.msgdef, rconn.msgtype))
register_types(typs)
for rconn, timestamp, data in reader.messages():
for rconn, timestamp, data in reader.messages(connections=connections):
data = ros1_to_cdr(data, rconn.msgtype)
writer.write(connmap[rconn.id], timestamp, data)
def convert_2to1(src: Path, dst: Path) -> None:
def convert_2to1(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None:
"""Convert Rosbag2 to Rosbag1.
Args:
src: Rosbag2 path.
dst: Rosbag1 path.
exclude_topics: Topics to skip.
Raises:
ConverterError: If all connections are excluded.
"""
with Reader2(src) as reader, Writer1(dst) as writer:
connmap: dict[int, Connection] = {}
for rconn in reader.connections:
connections = [x for x in reader.connections if x.topic not in exclude_topics]
if not connections:
raise ConverterError('No connections left for conversion.')
for rconn in connections:
candidate = downgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag1)
for conn in writer.connections:
@ -168,17 +181,18 @@ def convert_2to1(src: Path, dst: Path) -> None:
)
connmap[rconn.id] = conn
for rconn, timestamp, data in reader.messages():
for rconn, timestamp, data in reader.messages(connections=connections):
data = cdr_to_ros1(data, rconn.msgtype)
writer.write(connmap[rconn.id], timestamp, data)
def convert(src: Path, dst: Optional[Path]) -> None:
def convert(src: Path, dst: Optional[Path], exclude_topics: Sequence[str] = ()) -> None:
"""Convert between Rosbag1 and Rosbag2.
Args:
src: Source rosbag.
dst: Destination rosbag.
exclude_topics: Topics to skip.
Raises:
ConverterError: An error occured during reading, writing, or
@ -192,7 +206,7 @@ def convert(src: Path, dst: Optional[Path]) -> None:
func = convert_1to2 if upgrade else convert_2to1
try:
func(src, dst)
func(src, dst, exclude_topics)
except (ReaderError1, ReaderError2) as err:
raise ConverterError(f'Reading source bag: {err}') from err
except (WriterError1, WriterError2) as err:

View File

@ -42,7 +42,7 @@ def test_cliwrapper(tmp_path: Path) -> None:
with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]):
main()
cvrt.assert_called_with(tmp_path / 'ros1.bag', None)
cvrt.assert_called_with(src=tmp_path / 'ros1.bag', dst=None, exclude_topics=[])
with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt',
@ -68,7 +68,7 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--dst',
str(tmp_path / 'target')]):
main()
cvrt.assert_called_with(tmp_path / 'ros1.bag', tmp_path / 'target')
cvrt.assert_called_with(src=tmp_path / 'ros1.bag', dst=tmp_path / 'target', exclude_topics=[])
with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]), \
patch('builtins.print') as mock_print, \
@ -80,7 +80,7 @@ def test_cliwrapper(tmp_path: Path) -> None:
with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]):
main()
cvrt.assert_called_with(tmp_path / 'subdir', None)
cvrt.assert_called_with(src=tmp_path / 'subdir', dst=None, exclude_topics=[])
with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt',
@ -97,7 +97,7 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--dst',
str(tmp_path / 'target.bag')]):
main()
cvrt.assert_called_with(tmp_path / 'subdir', tmp_path / 'target.bag')
cvrt.assert_called_with(src=tmp_path / 'subdir', dst=tmp_path / 'target.bag', exclude_topics=[])
with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]), \
patch('builtins.print') as mock_print, \
@ -106,6 +106,14 @@ def test_cliwrapper(tmp_path: Path) -> None:
main()
mock_print.assert_called_with('ERROR: exc')
with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt',
str(tmp_path / 'ros1.bag'),
'--exclude-topic',
'/foo']):
main()
cvrt.assert_called_with(src=tmp_path / 'ros1.bag', dst=None, exclude_topics=['/foo'])
def test_convert_1to2(tmp_path: Path) -> None:
"""Test conversion from rosbag1 to rosbag2."""
@ -176,7 +184,7 @@ def test_convert_1to2(tmp_path: Path) -> None:
convert(Path('foo.bag'), None)
reader.assert_called_with(Path('foo.bag'))
readerinst.messages.assert_called_with()
readerinst.messages.assert_called_with(connections=readerinst.connections)
writer.assert_called_with(Path('foo'))
writerinst.add_connection.assert_has_calls(
@ -205,6 +213,9 @@ def test_convert_1to2(tmp_path: Path) -> None:
],
)
with pytest.raises(ConverterError, match='No connections left for conversion'):
convert(Path('foo.bag'), None, ['/topic', '/other'])
writerinst.connections.clear()
ros1_to_cdr.side_effect = KeyError('exc')
with pytest.raises(ConverterError, match='Converting rosbag: .*exc'):
@ -340,7 +351,9 @@ def test_convert_2to1(tmp_path: Path) -> None:
convert(Path('foo'), None)
reader.assert_called_with(Path('foo'))
reader.return_value.__enter__.return_value.messages.assert_called_with()
reader.return_value.__enter__.return_value.messages.assert_called_with(
connections=readerinst.connections,
)
writer.assert_called_with(Path('foo.bag'))
writer.return_value.__enter__.return_value.add_connection.assert_has_calls(
@ -389,6 +402,9 @@ def test_convert_2to1(tmp_path: Path) -> None:
],
)
with pytest.raises(ConverterError, match='No connections left for conversion'):
convert(Path('foobag'), None, ['/topic', '/other'])
writerinst.connections.clear()
cdr_to_ros1.side_effect = KeyError('exc')
with pytest.raises(ConverterError, match='Converting rosbag: .*exc'):