regen: specify any list of sources

ooh this is pretty nice
Edison-CBS · Jul 26, 2024 · ceb0b4a · ceb0b4a
1 parent 900d68b
commit ceb0b4a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/selfdrive/test/process_replay/regen.py b/selfdrive/test/process_replay/regen.py
@@ -6,6 +6,7 @@
 import numpy as np
 
 from typing import Any
+from functools import partial
 from collections.abc import Iterable
 
 from openpilot.selfdrive.test.process_replay.process_replay import CONFIGS, FAKEDATA, ProcessConfig, replay_process, get_process_config, \
@@ -14,7 +15,7 @@
 from openpilot.selfdrive.test.update_ci_routes import upload_route
 from openpilot.tools.lib.route import Route
 from openpilot.tools.lib.framereader import FrameReader, BaseFrameReader, FrameType
-from openpilot.tools.lib.logreader import LogReader, LogIterable
+from openpilot.tools.lib.logreader import LogReader, LogIterable, auto_source, internal_source, internal_source_zst
 from openpilot.tools.lib.helpers import save_log
 
 
@@ -75,7 +76,7 @@ def setup_data_readers(
         assert device_type != "neo", "Driver camera not supported on neo segments. Use dummy dcamera."
         frs['driverCameraState'] = FrameReader(r.dcamera_paths()[sidx])
   else:
-    lr = LogReader(f"cd:/{route.replace('|', '/')}/{sidx}/rlog.bz2")
+    lr = LogReader(f"{route}/{sidx}/r", default_source=partial(auto_source, sources=[internal_source, internal_source_zst]))
     frs = {}
     if needs_road_cam:
       frs['roadCameraState'] = FrameReader(f"cd:/{route.replace('|', '/')}/{sidx}/fcamera.hevc")

diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
@@ -179,23 +179,24 @@ def check_source(source: Source, *args) -> LogPaths:
   return files
 
 
-def auto_source(sr: SegmentRange, mode=ReadMode.RLOG) -> LogPaths:
+def auto_source(sr: SegmentRange, mode=ReadMode.RLOG, sources: list[Source] = None) -> LogPaths:
   if mode == ReadMode.SANITIZED:
     return comma_car_segments_source(sr, mode)
 
-  SOURCES: list[Source] = [internal_source, internal_source_zst, openpilotci_source, comma_api_source, comma_car_segments_source,]
+  if sources is None:
+    sources: list[Source] = [internal_source, internal_source_zst, openpilotci_source, comma_api_source, comma_car_segments_source,]
   exceptions = {}
 
   # for automatic fallback modes, auto_source needs to first check if rlogs exist for any source
   if mode in [ReadMode.AUTO, ReadMode.AUTO_INTERACTIVE]:
-    for source in SOURCES:
+    for source in sources:
       try:
         return check_source(source, sr, ReadMode.RLOG)
       except Exception:
         pass
 
   # Automatically determine viable source
-  for source in SOURCES:
+  for source in sources:
     try:
       return check_source(source, sr, mode)
     except Exception as e:
@@ -260,7 +261,7 @@ def _parse_identifiers(self, identifier: str | list[str]):
   def __init__(self, identifier: str | list[str], default_mode: ReadMode = ReadMode.RLOG,
                default_source=auto_source, sort_by_time=False, only_union_types=False):
     self.default_mode = default_mode
-    self.default_source = default_source
+    self.default_source = default_source  # TODO: `default_source` name is confusing? should be `source` and should not be able to be overridden
     self.identifier = identifier
 
     self.sort_by_time = sort_by_time