From 1309d42b64b3b019963b1e7b193e66de1b000f50 Mon Sep 17 00:00:00 2001 From: Marko Durkovic Date: Fri, 23 Sep 2022 12:35:37 +0200 Subject: [PATCH] Add include filters to rosbag conversion --- src/rosbags/convert/__main__.py | 12 +++++++-- src/rosbags/convert/converter.py | 42 +++++++++++++++++++++++++------- tests/test_convert.py | 35 ++++++++++++++++++++++---- 3 files changed, 73 insertions(+), 16 deletions(-) diff --git a/src/rosbags/convert/__main__.py b/src/rosbags/convert/__main__.py index dcd09ce2..07e62317 100644 --- a/src/rosbags/convert/__main__.py +++ b/src/rosbags/convert/__main__.py @@ -50,12 +50,20 @@ def main() -> None: type=pathtype(exists=False), help='destination path for converted rosbag', ) - parser.add_argument( + topic_group = parser.add_argument_group('filtering').add_mutually_exclusive_group() + topic_group.add_argument( '--exclude-topic', action='append', default=[], dest='exclude_topics', - help='exclude topic by name', + help='topic to exclude from conversion, even if included explicitly', + ) + topic_group.add_argument( + '--include-topic', + action='append', + default=[], + dest='include_topics', + help='topic to include in conversion, instead of all', ) args = parser.parse_args() diff --git a/src/rosbags/convert/converter.py b/src/rosbags/convert/converter.py index 4667541a..11e2e488 100644 --- a/src/rosbags/convert/converter.py +++ b/src/rosbags/convert/converter.py @@ -99,13 +99,19 @@ def downgrade_connection(rconn: Connection) -> Connection: ) -def convert_1to2(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: +def convert_1to2( + src: Path, + dst: Path, + exclude_topics: Sequence[str], + include_topics: Sequence[str], +) -> None: """Convert Rosbag1 to Rosbag2. Args: src: Rosbag1 path. dst: Rosbag2 path. - exclude_topics: Topics to skip. + exclude_topics: Topics to exclude from conversion, even if included explicitly. + include_topics: Topics to include in conversion, instead of all. Raises: ConverterError: If all connections are excluded. @@ -114,7 +120,10 @@ def convert_1to2(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: with Reader1(src) as reader, Writer2(dst) as writer: typs: dict[str, Any] = {} connmap: dict[int, Connection] = {} - connections = [x for x in reader.connections if x.topic not in exclude_topics] + connections = [ + x for x in reader.connections + if x.topic not in exclude_topics and (not include_topics or x.topic in include_topics) + ] if not connections: raise ConverterError('No connections left for conversion.') for rconn in connections: @@ -143,13 +152,19 @@ def convert_1to2(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: writer.write(connmap[rconn.id], timestamp, data) -def convert_2to1(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: +def convert_2to1( + src: Path, + dst: Path, + exclude_topics: Sequence[str], + include_topics: Sequence[str], +) -> None: """Convert Rosbag2 to Rosbag1. Args: src: Rosbag2 path. dst: Rosbag1 path. - exclude_topics: Topics to skip. + exclude_topics: Topics to exclude from conversion, even if included explicitly. + include_topics: Topics to include in conversion, instead of all. Raises: ConverterError: If all connections are excluded. @@ -157,7 +172,10 @@ def convert_2to1(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: """ with Reader2(src) as reader, Writer1(dst) as writer: connmap: dict[int, Connection] = {} - connections = [x for x in reader.connections if x.topic not in exclude_topics] + connections = [ + x for x in reader.connections + if x.topic not in exclude_topics and (not include_topics or x.topic in include_topics) + ] if not connections: raise ConverterError('No connections left for conversion.') for rconn in connections: @@ -186,13 +204,19 @@ def convert_2to1(src: Path, dst: Path, exclude_topics: Sequence[str]) -> None: writer.write(connmap[rconn.id], timestamp, data) -def convert(src: Path, dst: Optional[Path], exclude_topics: Sequence[str] = ()) -> None: +def convert( + src: Path, + dst: Optional[Path], + exclude_topics: Sequence[str] = (), + include_topics: Sequence[str] = (), +) -> None: """Convert between Rosbag1 and Rosbag2. Args: src: Source rosbag. dst: Destination rosbag. - exclude_topics: Topics to skip. + exclude_topics: Topics to exclude from conversion, even if included explicitly. + include_topics: Topics to include in conversion, instead of all. Raises: ConverterError: An error occured during reading, writing, or @@ -206,7 +230,7 @@ def convert(src: Path, dst: Optional[Path], exclude_topics: Sequence[str] = ()) func = convert_1to2 if upgrade else convert_2to1 try: - func(src, dst, exclude_topics) + func(src, dst, exclude_topics, include_topics) except (ReaderError1, ReaderError2) as err: raise ConverterError(f'Reading source bag: {err}') from err except (WriterError1, WriterError2) as err: diff --git a/tests/test_convert.py b/tests/test_convert.py index 77bb5fce..abe3bcb8 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -42,7 +42,12 @@ 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(src=tmp_path / 'ros1.bag', dst=None, exclude_topics=[]) + cvrt.assert_called_with( + src=tmp_path / 'ros1.bag', + dst=None, + exclude_topics=[], + include_topics=[], + ) with patch('rosbags.convert.__main__.convert') as cvrt, \ patch.object(sys, 'argv', ['cvt', @@ -68,7 +73,12 @@ def test_cliwrapper(tmp_path: Path) -> None: '--dst', str(tmp_path / 'target')]): main() - cvrt.assert_called_with(src=tmp_path / 'ros1.bag', dst=tmp_path / 'target', exclude_topics=[]) + cvrt.assert_called_with( + src=tmp_path / 'ros1.bag', + dst=tmp_path / 'target', + exclude_topics=[], + include_topics=[], + ) with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]), \ patch('builtins.print') as mock_print, \ @@ -80,7 +90,12 @@ 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(src=tmp_path / 'subdir', dst=None, exclude_topics=[]) + cvrt.assert_called_with( + src=tmp_path / 'subdir', + dst=None, + exclude_topics=[], + include_topics=[], + ) with patch('rosbags.convert.__main__.convert') as cvrt, \ patch.object(sys, 'argv', ['cvt', @@ -97,7 +112,12 @@ def test_cliwrapper(tmp_path: Path) -> None: '--dst', str(tmp_path / 'target.bag')]): main() - cvrt.assert_called_with(src=tmp_path / 'subdir', dst=tmp_path / 'target.bag', exclude_topics=[]) + cvrt.assert_called_with( + src=tmp_path / 'subdir', + dst=tmp_path / 'target.bag', + exclude_topics=[], + include_topics=[], + ) with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]), \ patch('builtins.print') as mock_print, \ @@ -112,7 +132,12 @@ def test_cliwrapper(tmp_path: Path) -> None: '--exclude-topic', '/foo']): main() - cvrt.assert_called_with(src=tmp_path / 'ros1.bag', dst=None, exclude_topics=['/foo']) + cvrt.assert_called_with( + src=tmp_path / 'ros1.bag', + dst=None, + exclude_topics=['/foo'], + include_topics=[], + ) def test_convert_1to2(tmp_path: Path) -> None: