move plotjuggler to segment range reader (#30969)

* pj

* limit test to 2 segs

* remove

* fix test

tools/plotjuggler/juggle.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
+import multiprocessing
 import os
 import sys
-import multiprocessing
 import platform
 import shutil
 import subprocess
@@ -9,13 +9,12 @@ import tarfile
 import tempfile
 import requests
 import argparse
+from functools import partial
 
 from openpilot.common.basedir import BASEDIR
-from openpilot.selfdrive.test.openpilotci import get_url
-from openpilot.tools.lib.logreader import LogReader
-from openpilot.tools.lib.route import Route, SegmentName
 from openpilot.tools.lib.helpers import save_log
-from urllib.parse import urlparse, parse_qs
+
+from openpilot.tools.lib.srreader import SegmentRangeReader
 
 juggle_dir = os.path.dirname(os.path.realpath(__file__))
 
@@ -53,17 +52,6 @@ def get_plotjuggler_version():
   return tuple(map(int, version.split(".")))
 
 
-def load_segment(segment_name):
-  if segment_name is None:
-    return []
-
-  try:
-    return list(LogReader(segment_name))
-  except (AssertionError, ValueError) as e:
-    print(f"Error parsing {segment_name}: {e}")
-    return []
-
-
 def start_juggler(fn=None, dbc=None, layout=None, route_or_segment_name=None):
   env = os.environ.copy()
   env["BASEDIR"] = BASEDIR
@@ -82,48 +70,16 @@ def start_juggler(fn=None, dbc=None, layout=None, route_or_segment_name=None):
   cmd = f'{PLOTJUGGLER_BIN} --buffer_size {MAX_STREAMING_BUFFER_SIZE} --plugin_folders {INSTALL_DIR}{extra_args}'
   subprocess.call(cmd, shell=True, env=env, cwd=juggle_dir)
 
+def process(can, lr):
+  return [d for d in lr if can or d.which() not in ['can', 'sendcan']]
 
-def juggle_route(route_or_segment_name, segment_count, qlog, can, layout, dbc=None, ci=False):
-  segment_start = 0
-  if 'cabana' in route_or_segment_name:
-    query = parse_qs(urlparse(route_or_segment_name).query)
-    route_or_segment_name = query["route"][0]
+def juggle_route(route_or_segment_name, can, layout, dbc=None):
+  sr = SegmentRangeReader(route_or_segment_name)
 
-  if route_or_segment_name.startswith(("http://", "https://", "cd:/")) or os.path.isfile(route_or_segment_name):
-    logs = [route_or_segment_name]
-  elif ci:
-    route_or_segment_name = SegmentName(route_or_segment_name, allow_route_name=True)
-    route = route_or_segment_name.route_name.canonical_name
-    segment_start = max(route_or_segment_name.segment_num, 0)
-    logs = [get_url(route, i) for i in range(100)]  # Assume there not more than 100 segments
-  else:
-    route_or_segment_name = SegmentName(route_or_segment_name, allow_route_name=True)
-    segment_start = max(route_or_segment_name.segment_num, 0)
-
-    if route_or_segment_name.segment_num != -1 and segment_count is None:
-      segment_count = 1
-
-    r = Route(route_or_segment_name.route_name.canonical_name, route_or_segment_name.data_dir)
-    logs = r.qlog_paths() if qlog else r.log_paths()
-
-  segment_end = segment_start + segment_count if segment_count else None
-  logs = logs[segment_start:segment_end]
-
-  if None in logs:
-    resp = input(f"{logs.count(None)}/{len(logs)} of the rlogs in this segment are missing, would you like to fall back to the qlogs? (y/n) ")
-    if resp == 'y':
-      logs = r.qlog_paths()[segment_start:segment_end]
-    else:
-      print("Please try a different route or segment")
-      return
-
-  all_data = []
   with multiprocessing.Pool(24) as pool:
-    for d in pool.map(load_segment, logs):
-      all_data += d
-
-  if not can:
-    all_data = [d for d in all_data if d.which() not in ['can', 'sendcan']]
+    all_data = []
+    for p in pool.map(partial(process, can), sr.lrs):
+      all_data.extend(p)
 
   # Infer DBC name from logs
   if dbc is None:
@@ -146,15 +102,12 @@ if __name__ == "__main__":
                                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
   parser.add_argument("--demo", action="store_true", help="Use the demo route instead of providing one")
-  parser.add_argument("--qlog", action="store_true", help="Use qlogs")
-  parser.add_argument("--ci", action="store_true", help="Download data from openpilot CI bucket")
   parser.add_argument("--can", action="store_true", help="Parse CAN data")
   parser.add_argument("--stream", action="store_true", help="Start PlotJuggler in streaming mode")
   parser.add_argument("--layout", nargs='?', help="Run PlotJuggler with a pre-defined layout")
   parser.add_argument("--install", action="store_true", help="Install or update PlotJuggler + plugins")
   parser.add_argument("--dbc", help="Set the DBC name to load for parsing CAN data. If not set, the DBC will be automatically inferred from the logs.")
   parser.add_argument("route_or_segment_name", nargs='?', help="The route or segment name to plot (cabana share URL accepted)")
-  parser.add_argument("segment_count", type=int, nargs='?', help="The number of segments to plot")
 
   if len(sys.argv) == 1:
     parser.print_help()
@@ -177,4 +130,4 @@ if __name__ == "__main__":
     start_juggler(layout=args.layout)
   else:
     route_or_segment_name = DEMO_ROUTE if args.demo else args.route_or_segment_name.strip()
-    juggle_route(route_or_segment_name, args.segment_count, args.qlog, args.can, args.layout, args.dbc, args.ci)
+    juggle_route(route_or_segment_name, args.can, args.layout, args.dbc)

tools/plotjuggler/test_plotjuggler.py

@@ -8,7 +8,7 @@ import unittest
 
 from openpilot.common.basedir import BASEDIR
 from openpilot.common.timeout import Timeout
-from openpilot.tools.plotjuggler.juggle import install
+from openpilot.tools.plotjuggler.juggle import DEMO_ROUTE, install
 
 PJ_DIR = os.path.join(BASEDIR, "tools/plotjuggler")
 
@@ -18,8 +18,8 @@ class TestPlotJuggler(unittest.TestCase):
     install()
 
     pj = os.path.join(PJ_DIR, "juggle.py")
-    with subprocess.Popen(f'QT_QPA_PLATFORM=offscreen {pj} --demo None 1 --qlog',
-                          stderr=subprocess.PIPE, shell=True, start_new_session=True) as p:
+    with subprocess.Popen(f'QT_QPA_PLATFORM=offscreen {pj} "{DEMO_ROUTE}/:2"',
+                           stderr=subprocess.PIPE, shell=True, start_new_session=True) as p:
       # Wait for "Done reading Rlog data" signal from the plugin
       output = "\n"
       with Timeout(180, error_msg=output):