diff --git a/py/farm_ng/core/events_file_reader.py b/py/farm_ng/core/events_file_reader.py index b1392559..a2a659ed 100644 --- a/py/farm_ng/core/events_file_reader.py +++ b/py/farm_ng/core/events_file_reader.py @@ -4,6 +4,7 @@ import json import struct import sys +from collections import defaultdict from dataclasses import dataclass from pathlib import Path from typing import IO, TYPE_CHECKING, Any, Generator, cast @@ -18,6 +19,7 @@ # public symbols __all__ = [ + "build_events_dict", "event_has_message", "EventsFileReader", "EventLogPosition", @@ -86,6 +88,23 @@ def payload_to_protobuf(event: Event, payload: bytes) -> Message: return message +def build_events_dict( + events_index: list[EventLogPosition], +) -> dict[str, list[EventLogPosition]]: + """Build a dictionary of lists of events, where the key is the path of the event. + + Args: + events_index (list[EventLogPosition]): List of events from a log file + + Returns: + dict[str, list[EventLogPosition]]: Dictionary of lists of events + """ + events_dict: dict[str, list[EventLogPosition]] = defaultdict(list) + for event_index in events_index: + events_dict[f"/{event_index.event.uri.path}"].append(event_index) + return events_dict + + @dataclass class EventLogPosition: """EventLogPosition is a dataclass that stores the event, position and reader.""" diff --git a/py/tests/test_core.py b/py/tests/test_core.py index eacffcc0..b00cace0 100644 --- a/py/tests/test_core.py +++ b/py/tests/test_core.py @@ -7,6 +7,7 @@ from farm_ng.core.events_file_reader import ( EventLogPosition, EventsFileReader, + build_events_dict, event_has_message, proto_from_json_file, ) @@ -267,6 +268,23 @@ def test_write_read(self, log_base: Path, reader_log_file: Path) -> None: assert reader.close() + def test_build_events_dict(self, log_base: Path, reader_log_file: Path) -> None: + num_events = 10 + with EventsFileWriter(file_base=log_base) as writer: + for i in range(num_events): + time_stamp = timestamp_pb2.Timestamp(stamp=i) + writer.write(path="hello", message=time_stamp) + writer.write(path="world", message=time_stamp) + writer.write(path="hello/world", message=time_stamp) + + with EventsFileReader(reader_log_file) as reader: + assert reader.is_open() + events_dict = build_events_dict(reader.get_index()) + assert len(events_dict) == 3 + assert len(events_dict["/hello"]) == num_events + assert len(events_dict["/world"]) == num_events + assert len(events_dict["/hello/world"]) == num_events + class TestEventsJson: def test_json_write_read_path(self, tmp_path: Path) -> None: