Add build_events_dict method (#166)

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."""

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: