Convert connections attribute to list

This commit is contained in:
Marko Durkovic
2022-04-21 15:15:10 +02:00
parent f7d69e35d5
commit 34ffe96692
14 changed files with 72 additions and 74 deletions
+4 -4
View File
@@ -111,10 +111,10 @@ def convert_1to2(src: Path, dst: Path) -> None:
typs: dict[str, Any] = {}
connmap: dict[int, Connection] = {}
for rconn in reader.connections.values():
for rconn in reader.connections:
candidate = upgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag2)
for conn in writer.connections.values():
for conn in writer.connections:
assert isinstance(conn.ext, ConnectionExtRosbag2)
if (
conn.topic == candidate.topic and conn.msgtype == candidate.msgtype and
@@ -147,10 +147,10 @@ def convert_2to1(src: Path, dst: Path) -> None:
"""
with Reader2(src) as reader, Writer1(dst) as writer:
connmap: dict[int, Connection] = {}
for rconn in reader.connections.values():
for rconn in reader.connections:
candidate = downgrade_connection(rconn)
assert isinstance(candidate.ext, ConnectionExtRosbag1)
for conn in writer.connections.values():
for conn in writer.connections:
assert isinstance(conn.ext, ConnectionExtRosbag1)
if (
conn.topic == candidate.topic and conn.md5sum == candidate.md5sum and
+2 -2
View File
@@ -143,11 +143,11 @@ class AnyReader:
typs: dict[str, Any] = {}
for reader in self.readers:
assert isinstance(reader, Reader1)
for connection in reader.connections.values():
for connection in reader.connections:
typs.update(get_types_from_msg(connection.msgdef, connection.msgtype))
register_types(typs, self.typestore)
self.connections = [y for x in self.readers for y in x.connections.values()]
self.connections = [y for x in self.readers for y in x.connections]
self.isopen = True
def close(self) -> None:
+17 -14
View File
@@ -342,7 +342,7 @@ class Reader:
raise ReaderError(f'File {str(self.path)!r} does not exist.')
self.bio: Optional[BinaryIO] = None
self.connections: dict[int, Connection] = {}
self.connections: list[Connection] = []
self.indexes: dict[int, list[IndexData]]
self.chunk_infos: list[ChunkInfo] = []
self.chunks: dict[int, Chunk] = {}
@@ -384,7 +384,7 @@ class Reader:
self.bio.seek(index_pos)
try:
self.connections = dict(self.read_connection() for _ in range(conn_count))
self.connections = [self.read_connection() for _ in range(conn_count)]
self.chunk_infos = [self.read_chunk_info() for _ in range(chunk_count)]
except ReaderError as err:
raise ReaderError(f'Bag index looks damaged: {err.args}') from None
@@ -402,14 +402,15 @@ class Reader:
self.indexes = {
cid: list(heapq.merge(*x, key=lambda x: x.time)) for cid, x in indexes.items()
}
assert all(self.indexes[x] for x in self.connections)
assert all(self.indexes[x.id] for x in self.connections)
for cid, connection in self.connections.items():
self.connections[cid] = Connection(
*connection[0:5],
len(self.indexes[cid]),
*connection[6:],
)
self.connections = [
Connection(
*x[0:5],
len(self.indexes[x.id]),
*x[6:],
) for x in self.connections
]
except ReaderError:
self.close()
raise
@@ -445,7 +446,7 @@ class Reader:
"""Topic information."""
topics = {}
for topic, group in groupby(
sorted(self.connections.values(), key=lambda x: x.topic),
sorted(self.connections, key=lambda x: x.topic),
key=lambda x: x.topic,
):
connections = list(group)
@@ -462,7 +463,7 @@ class Reader:
)
return topics
def read_connection(self) -> tuple[int, Connection]:
def read_connection(self) -> Connection:
"""Read connection record from current position."""
assert self.bio
header = Header.read(self.bio, RecordType.CONNECTION)
@@ -477,7 +478,7 @@ class Reader:
callerid = header.get_string('callerid') if 'callerid' in header else None
latching = int(header.get_string('latching')) if 'latching' in header else None
return conn, Connection(
return Connection(
conn,
topic,
normalize_msgtype(typ),
@@ -593,7 +594,9 @@ class Reader:
raise ReaderError('Rosbag is not open.')
if not connections:
connections = self.connections.values()
connections = self.connections
connmap = {x.id: x for x in self.connections}
indexes = [self.indexes[x.id] for x in connections]
for entry in heapq.merge(*indexes):
@@ -624,7 +627,7 @@ class Reader:
raise ReaderError('Expected to find message data.')
data = read_bytes(chunk, read_uint32(chunk))
connection = self.connections[header.get_uint32('conn')]
connection = connmap[header.get_uint32('conn')]
assert entry.time == header.get_time('time')
yield connection, entry.time, data
+5 -5
View File
@@ -159,7 +159,7 @@ class Writer:
self.bio: Optional[BinaryIO] = None
self.compressor: Callable[[bytes], bytes] = lambda x: x
self.compression_format = 'none'
self.connections: dict[int, Connection] = {}
self.connections: list[Connection] = []
self.chunks: list[WriteChunk] = [
WriteChunk(BytesIO(), -1, 2**64, 0, defaultdict(list)),
]
@@ -258,7 +258,7 @@ class Writer:
self,
)
if any(x[1:] == connection[1:] for x in self.connections.values()):
if any(x[1:] == connection[1:] for x in self.connections):
raise WriterError(
f'Connections can only be added once with same arguments: {connection!r}.',
)
@@ -266,7 +266,7 @@ class Writer:
bio = self.chunks[-1].data
self.write_connection(connection, bio)
self.connections[connection.id] = connection
self.connections.append(connection)
return connection
def write(self, connection: Connection, timestamp: int, data: bytes) -> None:
@@ -284,7 +284,7 @@ class Writer:
if not self.bio:
raise WriterError('Bag was not opened.')
if connection not in self.connections.values():
if connection not in self.connections:
raise WriterError(f'There is no connection {connection!r}.') from None
chunk = self.chunks[-1]
@@ -367,7 +367,7 @@ class Writer:
index_pos = self.bio.tell()
for connection in self.connections.values():
for connection in self.connections:
self.write_connection(connection, self.bio)
for chunk in self.chunks:
+6 -9
View File
@@ -136,8 +136,8 @@ class Reader:
if missing := [x for x in self.paths if not x.exists()]:
raise ReaderError(f'Some database files are missing: {[str(x) for x in missing]!r}')
self.connections = {
idx + 1: Connection(
self.connections = [
Connection(
id=idx + 1,
topic=x['topic_metadata']['name'],
msgtype=x['topic_metadata']['type'],
@@ -150,9 +150,9 @@ class Reader:
),
owner=self,
) for idx, x in enumerate(self.metadata['topics_with_message_count'])
}
]
noncdr = {
fmt for x in self.connections.values() if isinstance(x.ext, ConnectionExtRosbag2)
fmt for x in self.connections if isinstance(x.ext, ConnectionExtRosbag2)
if (fmt := x.ext.serialization_format) != 'cdr'
}
if noncdr:
@@ -209,10 +209,7 @@ class Reader:
@property
def topics(self) -> dict[str, TopicInfo]:
"""Topic information."""
return {
x.topic: TopicInfo(x.msgtype, x.msgdef, x.msgcount, [x])
for x in self.connections.values()
}
return {x.topic: TopicInfo(x.msgtype, x.msgdef, x.msgcount, [x]) for x in self.connections}
def messages( # pylint: disable=too-many-locals
self,
@@ -278,7 +275,7 @@ class Reader:
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]),
row[1]: next((x for x in self.connections if x.topic == row[0]),
None) # type: ignore
for row in cur
}
+6 -6
View File
@@ -81,7 +81,7 @@ class Writer: # pylint: disable=too-many-instance-attributes
self.compression_mode = ''
self.compression_format = ''
self.compressor: Optional[zstandard.ZstdCompressor] = None
self.connections: dict[int, Connection] = {}
self.connections: list[Connection] = []
self.counts: dict[int, int] = {}
self.conn: Optional[sqlite3.Connection] = None
self.cursor: Optional[sqlite3.Cursor] = None
@@ -152,7 +152,7 @@ class Writer: # pylint: disable=too-many-instance-attributes
raise WriterError('Bag was not opened.')
connection = Connection(
id=len(self.connections.values()) + 1,
id=len(self.connections) + 1,
topic=topic,
msgtype=msgtype,
msgdef='',
@@ -164,14 +164,14 @@ class Writer: # pylint: disable=too-many-instance-attributes
),
owner=self,
)
for conn in self.connections.values():
for conn in self.connections:
if (
conn.topic == connection.topic and conn.msgtype == connection.msgtype and
conn.ext == connection.ext
):
raise WriterError(f'Connection can only be added once: {connection!r}.')
self.connections[connection.id] = connection
self.connections.append(connection)
self.counts[connection.id] = 0
meta = (connection.id, topic, msgtype, serialization_format, offered_qos_profiles)
self.cursor.execute('INSERT INTO topics VALUES(?, ?, ?, ?, ?)', meta)
@@ -191,7 +191,7 @@ class Writer: # pylint: disable=too-many-instance-attributes
"""
if not self.cursor:
raise WriterError('Bag was not opened.')
if connection not in self.connections.values():
if connection not in self.connections:
raise WriterError(f'Tried to write to unknown connection {connection!r}.')
if self.compression_mode == 'message':
@@ -252,7 +252,7 @@ class Writer: # pylint: disable=too-many-instance-attributes
'offered_qos_profiles': x.ext.offered_qos_profiles,
},
'message_count': self.counts[x.id],
} for x in self.connections.values() if isinstance(x.ext, ConnectionExtRosbag2)
} for x in self.connections if isinstance(x.ext, ConnectionExtRosbag2)
],
'compression_format': self.compression_format,
'compression_mode': self.compression_mode,