From be2d5675b3ae42c16afbfa2647a70be385ab9aa8 Mon Sep 17 00:00:00 2001 From: Marko Durkovic Date: Mon, 13 Sep 2021 17:30:09 +0200 Subject: [PATCH] Fix connection mapping for reader2 messages --- src/rosbags/rosbag2/reader.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/rosbags/rosbag2/reader.py b/src/rosbags/rosbag2/reader.py index 636c0edc..e5045e88 100644 --- a/src/rosbags/rosbag2/reader.py +++ b/src/rosbags/rosbag2/reader.py @@ -17,7 +17,7 @@ from .connection import Connection if TYPE_CHECKING: from types import TracebackType - from typing import Any, Generator, Iterable, Literal, Optional, Type, Union + from typing import Any, Dict, Generator, Iterable, Literal, Optional, Type, Union class ReaderError(Exception): @@ -232,16 +232,23 @@ class Reader: if cur.fetchone()[0] != 2: raise ReaderError(f'Cannot open database {path} or database missing tables.') + cur.execute('SELECT name,id FROM topics') + connmap: Dict[int, Connection] = { + row[1]: next((x for x in self.connections.values() if x.topic == row[0]), + None) # type: ignore + for row in cur + } + cur.execute(querystr, args) if self.compression_mode == 'message': decomp = zstandard.ZstdDecompressor().decompress for row in cur: cid, timestamp, data = row - yield self.connections[cid], timestamp, decomp(data) + yield connmap[cid], timestamp, decomp(data) else: for cid, timestamp, data in cur: - yield self.connections[cid], timestamp, data + yield connmap[cid], timestamp, data def __enter__(self) -> Reader: """Open rosbag2 when entering contextmanager."""