Client: Graph: Add trips

prog/graph.py

@@ -8,6 +8,11 @@ from qgis.PyQt.QtWidgets import *
 from qgis.PyQt.QtCore import *
 from qgis.PyQt.QtGui import *
 
+import gtfs
+
+from qasync import QEventLoop, asyncClose, asyncSlot, QApplication
+import asyncio
+
 import numpy as np
 
 from matplotlib.backends.backend_qtagg import FigureCanvas
@@ -31,21 +36,37 @@ class Graph(QWidget):
 
     def __init__(self):
         super().__init__()
-        layout = QVBoxLayout()
-        self.setLayout(layout)
+        self._layout = QVBoxLayout()
+        self.setLayout(self._layout)
 
 
         canvas = FigureCanvas(Figure(figsize=(5, 3)))
-        layout.addWidget(canvas)
-        layout.addWidget(NavigationToolbar(canvas, self))
+        self._layout.addWidget(canvas)
+        self._layout.addWidget(NavigationToolbar(canvas, self))
 
         self.ax = canvas.figure.subplots()
 
 
 class HistoryGraph(Graph):
-    def __init__(self, stops):
+    def __init__(self, stops, default_line=None, default_dt_from=None, default_dt_to=None):
         super().__init__()
 
+        self._add_layout = QHBoxLayout()
+        self._layout.insertLayout(0, self._add_layout)
+
+        self._add_selector = QLineEdit(str(default_line) if default_line else "")
+        self._add_layout.addWidget(self._add_selector)
+        self._add_selector.returnPressed.connect(self.process_add)
+        self._add_dt_from = QLineEdit(format_dt_to_user(default_dt_from or datetime.datetime.now(local_timezone)-datetime.timedelta(hours=2)))
+        self._add_layout.addWidget(self._add_dt_from)
+        self._add_dt_from.returnPressed.connect(self.process_add)
+        self._add_dt_to = QLineEdit(format_dt_to_user(default_dt_to or datetime.datetime.now(local_timezone)))
+        self._add_layout.addWidget(self._add_dt_to)
+        self._add_dt_to.returnPressed.connect(self.process_add)
+        self._add_button = QPushButton("Add")
+        self._add_layout.addWidget(self._add_button)
+        self._add_button.clicked.connect(self.process_add)
+
         xticks = []
         xlabels = []
         for stop in stops:
@@ -97,3 +118,48 @@ class HistoryGraph(Graph):
 
         self._line.figure.canvas.draw()
 
+    @asyncSlot()
+    async def process_add(self):
+        selector = self._add_selector.text()
+        dt_from = parse_dt_from_user(self._add_dt_from.text())
+        dt_to = parse_dt_from_user(self._add_dt_to.text())
+
+        g = gtfs.for_date(dt_from)
+        await g.load_trips()
+
+        trips = [Trip(id, g.date) for id, t in g.trips.items() if id.startswith(selector+"_")]
+        ths = {t.trip_id: TripHistory(t) for t in trips}
+
+        route_ids = {trip.trip_id.split("_")[0] for trip in trips}
+        print(route_ids)
+
+        c = await get_communication()
+
+        for route_id in route_ids:
+            dt = dt_from.replace(second=0, microsecond=0, minute=0)
+            while dt < dt_to:
+                data, source_timestamps = await c.get_preprocessed_data(dt, route_id)
+                print("P1")
+                data = await unzip_parse(data)
+                print("P3")
+
+                for trip_id, th in ths.items():
+                    if trip_id in data:
+                        th.add_preprocessed_data(data[trip_id])
+                dt += datetime.timedelta(hours=1)
+
+            print("P4")
+
+        for route_id in route_ids:
+            stops_data = await c.gtfs_get_stop_times(dt_from, route_filter=route_id)
+
+            for trip_id, th in ths.items():
+                if trip_id.split("_")[0] == route_id:
+                    await th.load_stops(stops_data)
+
+        for th in ths.values():
+            if len(th.history):
+                print("ADD TRIP")
+                self.add_trip(th)
+        print("DONE")
+