diff --git a/tools/lib/route.py b/tools/lib/route.py index d3df93ccda..93788e7ab3 100644 --- a/tools/lib/route.py +++ b/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}" diff --git a/tools/lib/srreader.py b/tools/lib/srreader.py index 2c08072e98..08abd03d33 100644 --- a/tools/lib/srreader.py +++ b/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: diff --git a/tools/lib/tests/test_srreader.py b/tools/lib/tests/test_srreader.py index 84521ae219..e1c54557be 100644 --- a/tools/lib/tests/test_srreader.py +++ b/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)