Add include filters to rosbag conversion

This commit is contained in:
Marko Durkovic 2022-09-23 12:35:37 +02:00
parent 4437512f05
commit 1309d42b64
3 changed files with 73 additions and 16 deletions

View File

@ -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()

View File

@ -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:

View File

@ -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: