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), type=pathtype(exists=False),
help='destination path for converted rosbag', 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() args = parser.parse_args()
if args.dst is not None and (args.src.suffix == '.bag') == (args.dst.suffix == '.bag'): 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) sys.exit(1)
try: try:
convert(args.src, args.dst) convert(**args.__dict__)
except ConverterError as err: except ConverterError as err:
print(f'ERROR: {err}') # noqa: T001 print(f'ERROR: {err}') # noqa: T001
sys.exit(1) sys.exit(1)

View File

@ -21,7 +21,7 @@ from rosbags.typesys.msg import generate_msgdef
if TYPE_CHECKING: if TYPE_CHECKING:
from pathlib import Path from pathlib import Path
from typing import Any, Optional from typing import Any, Optional, Sequence
LATCH = """ LATCH = """
- history: 3 - 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. """Convert Rosbag1 to Rosbag2.
Args: Args:
src: Rosbag1 path. src: Rosbag1 path.
dst: Rosbag2 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: 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]
for rconn in reader.connections: if not connections:
raise ConverterError('No connections left for conversion.')
for rconn in connections:
candidate = upgrade_connection(rconn) candidate = upgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag2) assert isinstance(candidate.ext, ConnectionExtRosbag2)
for conn in writer.connections: 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)) typs.update(get_types_from_msg(rconn.msgdef, rconn.msgtype))
register_types(typs) 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) data = ros1_to_cdr(data, rconn.msgtype)
writer.write(connmap[rconn.id], timestamp, data) 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. """Convert Rosbag2 to Rosbag1.
Args: Args:
src: Rosbag2 path. src: Rosbag2 path.
dst: Rosbag1 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: with Reader2(src) as reader, Writer1(dst) as writer:
connmap: dict[int, Connection] = {} 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) candidate = downgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag1) assert isinstance(candidate.ext, ConnectionExtRosbag1)
for conn in writer.connections: for conn in writer.connections:
@ -168,17 +181,18 @@ def convert_2to1(src: Path, dst: Path) -> None:
) )
connmap[rconn.id] = conn 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) data = cdr_to_ros1(data, rconn.msgtype)
writer.write(connmap[rconn.id], timestamp, data) 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. """Convert between Rosbag1 and Rosbag2.
Args: Args:
src: Source rosbag. src: Source rosbag.
dst: Destination rosbag. dst: Destination rosbag.
exclude_topics: Topics to skip.
Raises: Raises:
ConverterError: An error occured during reading, writing, or 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 func = convert_1to2 if upgrade else convert_2to1
try: try:
func(src, dst) func(src, dst, exclude_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,7 @@ 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(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, \ with patch('rosbags.convert.__main__.convert') as cvrt, \
patch.object(sys, 'argv', ['cvt', patch.object(sys, 'argv', ['cvt',
@ -68,7 +68,7 @@ def test_cliwrapper(tmp_path: Path) -> None:
'--dst', '--dst',
str(tmp_path / 'target')]): str(tmp_path / 'target')]):
main() 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')]), \ 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 +80,7 @@ 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(tmp_path / 'subdir', None) cvrt.assert_called_with(src=tmp_path / 'subdir', dst=None, exclude_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 +97,7 @@ 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(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')]), \ with patch.object(sys, 'argv', ['cvt', str(tmp_path / 'subdir')]), \
patch('builtins.print') as mock_print, \ patch('builtins.print') as mock_print, \
@ -106,6 +106,14 @@ def test_cliwrapper(tmp_path: Path) -> None:
main() main()
mock_print.assert_called_with('ERROR: exc') 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: def test_convert_1to2(tmp_path: Path) -> None:
"""Test conversion from rosbag1 to rosbag2.""" """Test conversion from rosbag1 to rosbag2."""
@ -176,7 +184,7 @@ def test_convert_1to2(tmp_path: Path) -> None:
convert(Path('foo.bag'), None) convert(Path('foo.bag'), None)
reader.assert_called_with(Path('foo.bag')) 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')) writer.assert_called_with(Path('foo'))
writerinst.add_connection.assert_has_calls( 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() writerinst.connections.clear()
ros1_to_cdr.side_effect = KeyError('exc') ros1_to_cdr.side_effect = KeyError('exc')
with pytest.raises(ConverterError, match='Converting rosbag: .*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) convert(Path('foo'), None)
reader.assert_called_with(Path('foo')) 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.assert_called_with(Path('foo.bag'))
writer.return_value.__enter__.return_value.add_connection.assert_has_calls( 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() writerinst.connections.clear()
cdr_to_ros1.side_effect = KeyError('exc') cdr_to_ros1.side_effect = KeyError('exc')
with pytest.raises(ConverterError, match='Converting rosbag: .*exc'): with pytest.raises(ConverterError, match='Converting rosbag: .*exc'):