From 8768be968b47c4883c7f1eb6dd8f5749c22f65a8 Mon Sep 17 00:00:00 2001
From: Justin Newberry <justin@comma.ai>
Date: Tue, 23 Jan 2024 14:51:45 -0800
Subject: [PATCH] LogReader: don't use route for max_segment (#31129)

* dont use route

* fix and cleanup
old-commit-hash: d574013d7bda94a43b29095887e176b9723d13cb
---
 tools/lib/logreader.py            | 35 +++++++++++++++++++------------
 tools/lib/route.py                | 13 ++++++++++++
 tools/lib/tests/test_logreader.py |  5 ++---
 3 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
index e28f41341c..c29662dfbc 100755
--- a/tools/lib/logreader.py
+++ b/tools/lib/logreader.py
@@ -110,21 +110,31 @@ def apply_strategy(mode: ReadMode, rlog_paths, qlog_paths):
   elif mode == ReadMode.AUTO_INTERACIVE:
     return auto_strategy(rlog_paths, qlog_paths, True)
 
-def parse_slice(sr: SegmentRange, route: Route):
-  segs = np.arange(route.max_seg_number+1)
+def parse_slice(sr: SegmentRange):
   s = create_slice_from_string(sr._slice)
-  return segs[s] if isinstance(s, slice) else [segs[s]]
+  if isinstance(s, slice):
+    if s.stop is None or s.stop < 0 or (s.start is not None and s.start < 0): # we need the number of segments in order to parse this slice
+      segs = np.arange(sr.get_max_seg_number()+1)
+    else:
+      segs = np.arange(s.stop + 1)
+    return segs[s]
+  else:
+    if s < 0:
+      s = sr.get_max_seg_number() + s + 1
+    return [s]
+
+def comma_api_source(sr: SegmentRange, mode: ReadMode):
+  segs = parse_slice(sr)
 
-def comma_api_source(sr: SegmentRange, route: Route, mode: ReadMode):
-  segs = parse_slice(sr, route)
+  route = Route(sr.route_name)
 
   rlog_paths = [route.log_paths()[seg] for seg in segs]
   qlog_paths = [route.log_paths()[seg] for seg in segs]
 
   return apply_strategy(mode, rlog_paths, qlog_paths)
 
-def internal_source(sr: SegmentRange, route: Route, mode: ReadMode):
-  segs = parse_slice(sr, route)
+def internal_source(sr: SegmentRange, mode: ReadMode):
+  segs = parse_slice(sr)
 
   def get_internal_url(sr: SegmentRange, seg, file):
     return f"cd:/{sr.dongle_id}/{sr.timestamp}/{seg}/{file}.bz2"
@@ -134,16 +144,16 @@ def internal_source(sr: SegmentRange, route: Route, mode: ReadMode):
 
   return apply_strategy(mode, rlog_paths, qlog_paths)
 
-def openpilotci_source(sr: SegmentRange, route: Route, mode: ReadMode):
-  segs = parse_slice(sr, route)
+def openpilotci_source(sr: SegmentRange, mode: ReadMode):
+  segs = parse_slice(sr)
 
   rlog_paths = [get_url(sr.route_name, seg, "rlog") for seg in segs]
   qlog_paths = [get_url(sr.route_name, seg, "qlog")  for seg in segs]
 
   return apply_strategy(mode, rlog_paths, qlog_paths)
 
-def comma_car_segments_source(sr: SegmentRange, route: Route, mode=ReadMode.RLOG):
-  segs = parse_slice(sr, route)
+def comma_car_segments_source(sr: SegmentRange, mode=ReadMode.RLOG):
+  segs = parse_slice(sr)
 
   return [get_comma_segments_url(sr.route_name, seg) for seg in segs]
 
@@ -209,11 +219,10 @@ class LogReader:
         return direct_source(identifier)
 
     sr = SegmentRange(parsed)
-    route = Route(sr.route_name)
     mode = self.default_mode if sr.selector is None else ReadMode(sr.selector)
     source = self.default_source if source is None else source
 
-    return source(sr, route, mode)
+    return source(sr, mode)
 
   def __init__(self, identifier: str | List[str], default_mode=ReadMode.RLOG, default_source=auto_source, sort_by_time=False, only_union_types=False):
     self.default_mode = default_mode
diff --git a/tools/lib/route.py b/tools/lib/route.py
index 09eb7f5426..f7c3c432c8 100644
--- a/tools/lib/route.py
+++ b/tools/lib/route.py
@@ -1,5 +1,6 @@
 import os
 import re
+from functools import cache
 from urllib.parse import urlparse
 from collections import defaultdict
 from itertools import chain
@@ -231,11 +232,23 @@ class SegmentName:
   def __str__(self) -> str: return self._canonical_name
 
 
+@cache
+def get_max_seg_number_cached(sr: 'SegmentRange'):
+  try:
+    api = CommaApi(get_token())
+    return api.get("/v1/route/" + sr.route_name.replace("/", "|"))["segment_numbers"][-1]
+  except Exception as e:
+    raise Exception("unable to get max_segment_number. ensure you have access to this route or the route is public.") from e
+
+
 class SegmentRange:
   def __init__(self, segment_range: str):
     self.m = re.fullmatch(RE.SEGMENT_RANGE, segment_range)
     assert self.m, f"Segment range is not valid {segment_range}"
 
+  def get_max_seg_number(self):
+    return get_max_seg_number_cached(self)
+
   @property
   def route_name(self):
     return self.m.group("route_name")
diff --git a/tools/lib/tests/test_logreader.py b/tools/lib/tests/test_logreader.py
index 6ec3edc9ac..f7874a3fb3 100644
--- a/tools/lib/tests/test_logreader.py
+++ b/tools/lib/tests/test_logreader.py
@@ -6,7 +6,7 @@ import pytest
 from parameterized import parameterized
 import requests
 from openpilot.tools.lib.logreader import LogReader, parse_indirect, parse_slice, ReadMode
-from openpilot.tools.lib.route import Route, SegmentRange
+from openpilot.tools.lib.route import SegmentRange
 
 NUM_SEGS = 17 # number of segments in the test route
 ALL_SEGS = list(np.arange(NUM_SEGS))
@@ -42,8 +42,7 @@ class TestLogReader(unittest.TestCase):
   def test_indirect_parsing(self, identifier, expected):
     parsed, _, _ = parse_indirect(identifier)
     sr = SegmentRange(parsed)
-    route = Route(sr.route_name)
-    segs = parse_slice(sr, route)
+    segs = parse_slice(sr)
     self.assertListEqual(list(segs), expected)
 
   @parameterized.expand([