Replace multilogiterator (#30980)

tools/lib/README.md

@@ -31,21 +31,3 @@ for msg in lr:
   if msg.which() == "carState":
     print(msg.carState.steeringAngleDeg)
 ```
-
-### MultiLogIterator
-
-`MultiLogIterator` is similar to `LogReader`, but reads multiple logs. 
-
-```python
-from openpilot.tools.lib.route import Route
-from openpilot.tools.lib.logreader import MultiLogIterator
-
-# setup a MultiLogIterator to read all the logs in the route
-r = Route("a2a0ccea32023010|2023-07-27--13-01-19")
-lr = MultiLogIterator(r.log_paths())
-
-# print all the steering angles values from all the logs in the route
-for msg in lr:
-  if msg.which() == "carState":
-    print(msg.carState.steeringAngleDeg)
-```

tools/lib/logreader.py

@@ -10,71 +10,9 @@ from typing import Iterable, Iterator
 
 from cereal import log as capnp_log
 from openpilot.tools.lib.filereader import FileReader
-from openpilot.tools.lib.route import Route, SegmentName
 
 LogIterable = Iterable[capnp._DynamicStructReader]
 
-# this is an iterator itself, and uses private variables from LogReader
-class MultiLogIterator:
-  def __init__(self, log_paths, sort_by_time=False):
-    self._log_paths = log_paths
-    self.sort_by_time = sort_by_time
-
-    self._first_log_idx = next(i for i in range(len(log_paths)) if log_paths[i] is not None)
-    self._current_log = self._first_log_idx
-    self._idx = 0
-    self._log_readers = [None]*len(log_paths)
-    self.start_time = self._log_reader(self._first_log_idx)._ts[0]
-
-  def _log_reader(self, i):
-    if self._log_readers[i] is None and self._log_paths[i] is not None:
-      log_path = self._log_paths[i]
-      self._log_readers[i] = LogReader(log_path, sort_by_time=self.sort_by_time)
-
-    return self._log_readers[i]
-
-  def __iter__(self) -> Iterator[capnp._DynamicStructReader]:
-    return self
-
-  def _inc(self):
-    lr = self._log_reader(self._current_log)
-    if self._idx < len(lr._ents)-1:
-      self._idx += 1
-    else:
-      self._idx = 0
-      self._current_log = next(i for i in range(self._current_log + 1, len(self._log_readers) + 1)
-                               if i == len(self._log_readers) or self._log_paths[i] is not None)
-      if self._current_log == len(self._log_readers):
-        raise StopIteration
-
-  def __next__(self):
-    while 1:
-      lr = self._log_reader(self._current_log)
-      ret = lr._ents[self._idx]
-      self._inc()
-      return ret
-
-  def tell(self):
-    # returns seconds from start of log
-    return (self._log_reader(self._current_log)._ts[self._idx] - self.start_time) * 1e-9
-
-  def seek(self, ts):
-    # seek to nearest minute
-    minute = int(ts/60)
-    if minute >= len(self._log_paths) or self._log_paths[minute] is None:
-      return False
-
-    self._current_log = minute
-
-    # HACK: O(n) seek afterward
-    self._idx = 0
-    while self.tell() < ts:
-      self._inc()
-    return True
-
-  def reset(self):
-    self.__init__(self._log_paths, sort_by_time=self.sort_by_time)
-
 
 class LogReader:
   def __init__(self, fn, canonicalize=True, only_union_types=False, sort_by_time=False, dat=None):
@@ -121,14 +59,6 @@ class LogReader:
       else:
         yield ent
 
-def logreader_from_route_or_segment(r, sort_by_time=False):
-  sn = SegmentName(r, allow_route_name=True)
-  route = Route(sn.route_name.canonical_name)
-  if sn.segment_num < 0:
-    return MultiLogIterator(route.log_paths(), sort_by_time=sort_by_time)
-  else:
-    return LogReader(route.log_paths()[sn.segment_num], sort_by_time=sort_by_time)
-
 
 if __name__ == "__main__":
   import codecs

tools/lib/srreader.py

@@ -1,4 +1,5 @@
 import enum
+import itertools
 import numpy as np
 import pathlib
 import re
@@ -132,10 +133,16 @@ class SegmentRangeReader:
     self.default_mode = default_mode
     self.default_source = default_source
     self.sort_by_time = sort_by_time
+    self.identifier = identifier
 
-    self.lrs = self._logreaders_from_identifier(identifier)
+    self.reset()
 
   def __iter__(self):
-    for lr in self.lrs:
-      for m in lr:
-        yield m
+    return self
+
+  def __next__(self):
+    return next(self.chain)
+
+  def reset(self):
+    self.lrs = self._logreaders_from_identifier(self.identifier)
+    self.chain = itertools.chain(*self.lrs)

tools/scripts/save_ubloxraw_stream.py

@@ -3,8 +3,7 @@ import argparse
 import os
 import sys
 from openpilot.common.basedir import BASEDIR
-from openpilot.tools.lib.logreader import MultiLogIterator
-from openpilot.tools.lib.route import Route
+from openpilot.tools.lib.srreader import SegmentRangeReader
 
 os.environ['BASEDIR'] = BASEDIR
 
@@ -14,28 +13,17 @@ def get_arg_parser():
       description="Unlogging and save to file",
       formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-  parser.add_argument("data_dir", nargs='?',
-                              help="Path to directory in which log and camera files are located.")
-  parser.add_argument("route_name", type=(lambda x: x.replace("#", "|")), nargs="?",
+  parser.add_argument("route", type=(lambda x: x.replace("#", "|")), nargs="?",
                       help="The route whose messages will be published.")
   parser.add_argument("--out_path", nargs='?', default='/data/ubloxRaw.stream',
                       help="Output pickle file path")
   return parser
 
 
-def main(argv):
+def main():
   args = get_arg_parser().parse_args(sys.argv[1:])
-  if not args.data_dir:
-    print('Data directory invalid.')
-    return
 
-  if not args.route_name:
-    # Extract route name from path
-    args.route_name = os.path.basename(args.data_dir)
-    args.data_dir = os.path.dirname(args.data_dir)
-
-  route = Route(args.route_name, args.data_dir)
-  lr = MultiLogIterator(route.log_paths())
+  lr = SegmentRangeReader(args.route)
 
   with open(args.out_path, 'wb') as f:
     try:
@@ -56,4 +44,4 @@ def main(argv):
 
 
 if __name__ == "__main__":
-  sys.exit(main(sys.argv[1:]))
+  main()