segmentrangereader: support more sources (#30970)

* support more sources

* remove this
old-commit-hash: fb994ae782259e8ed4d9212edc3c9fb277f831f6

tools/lib/route.py

@@ -1,6 +1,6 @@
 import os
 import re
-from urllib.parse import urlparse
+from urllib.parse import parse_qs, urlparse
 from collections import defaultdict
 from itertools import chain
 from typing import Optional
@@ -231,8 +231,27 @@ class SegmentName:
   def __str__(self) -> str: return self._canonical_name
 
 
+def parse_useradmin(segment_range):
+  if "useradmin.comma.ai" in segment_range:
+    query = parse_qs(urlparse(segment_range).query)
+    return query["onebox"][0]
+  return segment_range
+
+def parse_cabana(segment_range):
+  if "cabana.comma.ai" in segment_range:
+    query = parse_qs(urlparse(segment_range).query)
+    return query["route"][0]
+  return segment_range
+
+def parse_cd(segment_range):
+  return segment_range.replace("cd:/", "")
+
 class SegmentRange:
   def __init__(self, segment_range: str):
+    segment_range = parse_useradmin(segment_range)
+    segment_range = parse_cabana(segment_range)
+    segment_range = parse_cd(segment_range)
+
     self.m = re.fullmatch(RE.SEGMENT_RANGE, segment_range)
     assert self.m, f"Segment range is not valid {segment_range}"
 

tools/lib/srreader.py

@@ -71,12 +71,12 @@ def auto_source(*args, **kwargs):
 
 
 class SegmentRangeReader:
-  def __init__(self, segment_range: str, default_mode=ReadMode.RLOG, source=auto_source, sort_by_time=False):
+  def __init__(self, segment_range: str, default_mode=ReadMode.RLOG, default_source=auto_source, sort_by_time=False):
     sr = SegmentRange(segment_range)
 
     mode = default_mode if sr.selector is None else ReadMode(sr.selector)
 
-    self.lrs = source(sr, mode, sort_by_time)
+    self.lrs = default_source(sr, mode, sort_by_time)
 
   def __iter__(self):
     for lr in self.lrs:

tools/lib/tests/test_srreader.py

@@ -30,6 +30,11 @@ class TestSegmentRangeReader(unittest.TestCase):
     (f"{TEST_ROUTE}/-4:-2", ALL_SEGS[-4:-2]),
     (f"{TEST_ROUTE}/:10:2", ALL_SEGS[:10:2]),
     (f"{TEST_ROUTE}/5::2", ALL_SEGS[5::2]),
+    (f"https://useradmin.comma.ai/?onebox={TEST_ROUTE}", ALL_SEGS),
+    (f"https://useradmin.comma.ai/?onebox={TEST_ROUTE.replace('/', '|')}", ALL_SEGS),
+    (f"https://useradmin.comma.ai/?onebox={TEST_ROUTE.replace('/', '%7C')}", ALL_SEGS),
+    (f"https://cabana.comma.ai/?route={TEST_ROUTE}", ALL_SEGS),
+    (f"cd:/{TEST_ROUTE}", ALL_SEGS),
   ])
   def test_parse_slice(self, segment_range, expected):
     sr = SegmentRange(segment_range)