Unify rosbag1 and rosbag2 topic information

This commit is contained in:
Marko Durkovic 2022-04-13 10:42:47 +02:00
parent 16d1758327
commit d32742b904
3 changed files with 39 additions and 43 deletions

View File

@ -34,3 +34,12 @@ class Connection(NamedTuple):
md5sum: str md5sum: str
msgcount: int msgcount: int
ext: Union[ConnectionExtRosbag1, ConnectionExtRosbag2] ext: Union[ConnectionExtRosbag1, ConnectionExtRosbag2]
class TopicInfo(NamedTuple):
"""Topic information."""
msgtype: Optional[str]
msgdef: Optional[str]
msgcount: int
connections: list[Connection]

View File

@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Dict, NamedTuple
from lz4.frame import decompress as lz4_decompress from lz4.frame import decompress as lz4_decompress
from rosbags.interfaces import Connection, ConnectionExtRosbag1 from rosbags.interfaces import Connection, ConnectionExtRosbag1, TopicInfo
from rosbags.typesys.msg import normalize_msgtype from rosbags.typesys.msg import normalize_msgtype
if TYPE_CHECKING: if TYPE_CHECKING:
@ -68,15 +68,6 @@ class Chunk(NamedTuple):
decompressor: Callable[[bytes], bytes] decompressor: Callable[[bytes], bytes]
class TopicInfo(NamedTuple):
"""Topic information."""
conncount: int
msgcount: int
msgdef: str
msgtype: str
class IndexData(NamedTuple): class IndexData(NamedTuple):
"""Index data.""" """Index data."""
@ -336,8 +327,6 @@ class Reader:
""" """
# pylint: disable=too-many-instance-attributes
def __init__(self, path: Union[str, Path]): def __init__(self, path: Union[str, Path]):
"""Initialize. """Initialize.
@ -358,7 +347,6 @@ class Reader:
self.chunk_infos: list[ChunkInfo] = [] self.chunk_infos: list[ChunkInfo] = []
self.chunks: dict[int, Chunk] = {} self.chunks: dict[int, Chunk] = {}
self.current_chunk: tuple[int, BinaryIO] = (-1, BytesIO()) self.current_chunk: tuple[int, BinaryIO] = (-1, BytesIO())
self.topics: dict[str, TopicInfo] = {}
def open(self) -> None: # pylint: disable=too-many-locals def open(self) -> None: # pylint: disable=too-many-locals
"""Open rosbag and read metadata.""" """Open rosbag and read metadata."""
@ -422,28 +410,6 @@ class Reader:
len(self.indexes[cid]), len(self.indexes[cid]),
connection[6], connection[6],
) )
self.topics = {}
for topic, group in groupby(
sorted(self.connections.values(), key=lambda x: x.topic),
key=lambda x: x.topic,
):
connections = list(group)
count = reduce(
lambda x, y: x + y,
(
y.connection_counts.get(x.id, 0)
for x in connections
for y in self.chunk_infos
),
)
self.topics[topic] = TopicInfo(
len(connections),
count,
connections[0].msgdef,
connections[0].msgtype,
)
except ReaderError: except ReaderError:
self.close() self.close()
raise raise
@ -474,6 +440,28 @@ class Reader:
"""Total message count.""" """Total message count."""
return reduce(lambda x, y: x + y, (x.msgcount for x in self.topics.values()), 0) return reduce(lambda x, y: x + y, (x.msgcount for x in self.topics.values()), 0)
@property
def topics(self) -> dict[str, TopicInfo]:
"""Topic information."""
topics = {}
for topic, group in groupby(
sorted(self.connections.values(), key=lambda x: x.topic),
key=lambda x: x.topic,
):
connections = list(group)
msgcount = reduce(
lambda x, y: x + y,
(y.connection_counts.get(x.id, 0) for x in connections for y in self.chunk_infos),
)
topics[topic] = TopicInfo(
msgtypes.pop() if len(msgtypes := {x.msgtype for x in connections}) == 1 else None,
msgdefs.pop() if len(msgdefs := {x.msgdef for x in connections}) == 1 else None,
msgcount,
connections,
)
return topics
def read_connection(self) -> tuple[int, Connection]: def read_connection(self) -> tuple[int, Connection]:
"""Read connection record from current position.""" """Read connection record from current position."""
assert self.bio assert self.bio

View File

@ -14,7 +14,7 @@ import zstandard
from ruamel.yaml import YAML from ruamel.yaml import YAML
from ruamel.yaml.error import YAMLError from ruamel.yaml.error import YAMLError
from rosbags.interfaces import Connection, ConnectionExtRosbag2 from rosbags.interfaces import Connection, ConnectionExtRosbag2, TopicInfo
if TYPE_CHECKING: if TYPE_CHECKING:
from types import TracebackType from types import TracebackType
@ -206,13 +206,12 @@ class Reader:
return mode if mode != 'none' else None return mode if mode != 'none' else None
@property @property
def topics(self) -> dict[str, Connection]: def topics(self) -> dict[str, TopicInfo]:
"""Topic information. """Topic information."""
return {
For the moment this a dictionary mapping topic names to connections. x.topic: TopicInfo(x.msgtype, x.msgdef, x.msgcount, [x])
for x in self.connections.values()
""" }
return {x.topic: x for x in self.connections.values()}
def messages( # pylint: disable=too-many-locals def messages( # pylint: disable=too-many-locals
self, self,