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), type=pathtype(exists=False),
help='destination path for converted rosbag', 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', '--exclude-topic',
action='append', action='append',
default=[], default=[],
dest='exclude_topics', 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() 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. """Convert Rosbag1 to Rosbag2.
Args: Args:
src: Rosbag1 path. src: Rosbag1 path.
dst: Rosbag2 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: Raises:
ConverterError: If all connections are excluded. 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: with Reader1(src) as reader, Writer2(dst) as writer:
typs: dict[str, Any] = {} typs: dict[str, Any] = {}
connmap: dict[int, Connection] = {} 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: if not connections:
raise ConverterError('No connections left for conversion.') raise ConverterError('No connections left for conversion.')
for rconn in connections: 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) 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. """Convert Rosbag2 to Rosbag1.
Args: Args:
src: Rosbag2 path. src: Rosbag2 path.
dst: Rosbag1 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: Raises:
ConverterError: If all connections are excluded. 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: with Reader2(src) as reader, Writer1(dst) as writer:
connmap: dict[int, Connection] = {} 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: if not connections:
raise ConverterError('No connections left for conversion.') raise ConverterError('No connections left for conversion.')
for rconn in connections: 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) 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. """Convert between Rosbag1 and Rosbag2.
Args: Args:
src: Source rosbag. src: Source rosbag.
dst: Destination 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: Raises:
ConverterError: An error occured during reading, writing, or 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 func = convert_1to2 if upgrade else convert_2to1
try: try:
func(src, dst, exclude_topics) func(src, dst, exclude_topics, include_topics)
except (ReaderError1, ReaderError2) as err: except (ReaderError1, ReaderError2) as err:
raise ConverterError(f'Reading source bag: {err}') from err raise ConverterError(f'Reading source bag: {err}') from err
except (WriterError1, WriterError2) as 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, \ with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]): patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]):
main() 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, \ with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', patch.object(sys, 'argv', ['cvt',
@ -68,7 +73,12 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--dst', '--dst',
str(tmp_path / 'target')]): str(tmp_path / 'target')]):
main() 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')]), \ with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'ros1.bag')]), \
patch('builtins.print') as mock_print, \ 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, \ with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]): patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]):
main() 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, \ with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', patch.object(sys, 'argv', ['cvt',
@ -97,7 +112,12 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--dst', '--dst',
str(tmp_path / 'target.bag')]): str(tmp_path / 'target.bag')]):
main() 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')]), \ with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]), \
patch('builtins.print') as mock_print, \ patch('builtins.print') as mock_print, \
@ -112,7 +132,12 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--exclude-topic', '--exclude-topic',
'/foo']): '/foo']):
main() 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: def test_convert_1to2(tmp_path: Path) -> None: