From 1eef956cadd30aa4fd5aca154ec7074f16a849df Mon Sep 17 00:00:00 2001
From: Shane Smiskol <shane@smiskol.com>
Date: Thu, 14 Aug 2025 19:19:37 -0700
Subject: [PATCH] LogReader sourcing: return dict (#35994)

* new return type

* fix test

* why not
---
 tools/lib/logreader.py            | 33 ++++++++++++++++---------------
 tools/lib/tests/test_logreader.py |  6 +++---
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
index 6b0d9fa527..d15e473c41 100755
--- a/tools/lib/logreader.py
+++ b/tools/lib/logreader.py
@@ -141,7 +141,7 @@ class ReadMode(enum.StrEnum):
 
 LogPath = str | None
 LogFileName = tuple[str, ...]
-Source = Callable[[SegmentRange, LogFileName], list[LogPath]]
+Source = Callable[[SegmentRange, LogFileName], dict[int, LogPath]]
 
 InternalUnavailableException = Exception("Internal source not available")
 
@@ -150,52 +150,53 @@ class LogsUnavailable(Exception):
   pass
 
 
-def comma_api_source(sr: SegmentRange, fns: LogFileName) -> list[LogPath]:
+def comma_api_source(sr: SegmentRange, fns: LogFileName) -> dict[int, LogPath]:
   route = Route(sr.route_name)
 
   # comma api will have already checked if the file exists
   if fns == FileName.RLOG:
-    return [route.log_paths()[seg] for seg in sr.seg_idxs]
+    return {seg: route.log_paths()[seg] for seg in sr.seg_idxs}
   else:
-    return [route.qlog_paths()[seg] for seg in sr.seg_idxs]
+    return {seg: route.qlog_paths()[seg] for seg in sr.seg_idxs}
 
 
-def internal_source(sr: SegmentRange, fns: LogFileName, endpoint_url: str = DATA_ENDPOINT) -> list[LogPath]:
+def internal_source(sr: SegmentRange, fns: LogFileName, endpoint_url: str = DATA_ENDPOINT) -> dict[int, LogPath]:
   if not internal_source_available(endpoint_url):
     raise InternalUnavailableException
 
   def get_internal_url(sr: SegmentRange, seg, file):
     return f"{endpoint_url.rstrip('/')}/{sr.dongle_id}/{sr.log_id}/{seg}/{file}"
 
-  return eval_source([[get_internal_url(sr, seg, fn) for fn in fns] for seg in sr.seg_idxs])
+  return eval_source({seg: [get_internal_url(sr, seg, fn) for fn in fns] for seg in sr.seg_idxs})
 
 
-def openpilotci_source(sr: SegmentRange, fns: LogFileName) -> list[LogPath]:
-  return eval_source([[get_url(sr.route_name, seg, fn) for fn in fns] for seg in sr.seg_idxs])
+def openpilotci_source(sr: SegmentRange, fns: LogFileName) -> dict[int, LogPath]:
+  return eval_source({seg: [get_url(sr.route_name, seg, fn) for fn in fns] for seg in sr.seg_idxs})
 
 
-def comma_car_segments_source(sr: SegmentRange, fns: LogFileName) -> list[LogPath]:
-  return eval_source([get_comma_segments_url(sr.route_name, seg) for seg in sr.seg_idxs])
+def comma_car_segments_source(sr: SegmentRange, fns: LogFileName) -> dict[int, LogPath]:
+  return eval_source({seg: get_comma_segments_url(sr.route_name, seg) for seg in sr.seg_idxs})
 
 
 def direct_source(file_or_url: str) -> list[str]:
   return [file_or_url]
 
 
-def eval_source(files: list[list[str] | str]) -> list[LogPath]:
+def eval_source(files: dict[int, list[str] | str]) -> dict[int, LogPath]:
   # Returns valid file URLs given a list of possible file URLs for each segment (e.g. rlog.bz2, rlog.zst)
-  valid_files: list[LogPath] = []
+  valid_files: dict[int, LogPath] = {}
 
-  for urls in files:
+  for seg_idx, urls in files.items():
     if isinstance(urls, str):
       urls = [urls]
 
+    # Add first valid file URL or None
     for url in urls:
       if file_exists(url):
-        valid_files.append(url)
+        valid_files[seg_idx] = url
         break
     else:
-      valid_files.append(None)
+      valid_files[seg_idx] = None
 
   return valid_files
 
@@ -227,7 +228,7 @@ def auto_source(identifier: str, sources: list[Source], default_mode: ReadMode)
         assert len(files) == len(valid_files) or len(valid_files) == 0, f"Source {source.__name__} returned unexpected number of files"
 
         # Build a dict of valid files
-        for idx, f in enumerate(files):
+        for idx, f in files.items():
           if valid_files.get(idx) is None:
             valid_files[idx] = f
 
diff --git a/tools/lib/tests/test_logreader.py b/tools/lib/tests/test_logreader.py
index 11bdd33ccf..7a70ad6aab 100644
--- a/tools/lib/tests/test_logreader.py
+++ b/tools/lib/tests/test_logreader.py
@@ -36,12 +36,12 @@ def setup_source_scenario(mocker, is_internal=False):
   comma_api_source_mock.__name__ = comma_api_source_mock._mock_name
 
   if is_internal:
-    internal_source_mock.return_value = [QLOG_FILE]
+    internal_source_mock.return_value = {3: QLOG_FILE}
   else:
     internal_source_mock.side_effect = InternalUnavailableException
 
-  openpilotci_source_mock.return_value = [None]
-  comma_api_source_mock.return_value = [QLOG_FILE]
+  openpilotci_source_mock.return_value = {3: None}
+  comma_api_source_mock.return_value = {3: QLOG_FILE}
 
   yield