add type to pivoted trains

robinklaassen · May 17, 2024 · 055266b · 055266b
1 parent 788d5c0
commit 055266b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/aid/provide/ns_trains.py b/aid/provide/ns_trains.py
@@ -15,12 +15,13 @@ class TrainRecord(BaseModel):
     speed: float
     direction: float
     accuracy: float
+    type: str
 
 
 class NSTrainProvider(BaseProvider):
     def get_trains(self, start: datetime, end: datetime) -> list[TrainRecord]:
         query = """
-        select timestamp, id, round(x) as x, round(y) as y, speed, direction, accuracy
+        select timestamp, id, round(x) as x, round(y) as y, speed, direction, accuracy, type
         from std.trains
         where timestamp between %s and %s
         order by timestamp asc, id asc

diff --git a/aid/provide/routers/trains.py b/aid/provide/routers/trains.py
@@ -88,7 +88,7 @@ def get_pivoted_data(start: datetime | None = None, end: datetime | None = None)
     df = df.round(5)
 
     # Pivot to requested format
-    df = df.melt(id_vars=["timestamp", "id"], value_vars=["x", "y", "speed"], var_name="var")
+    df = df.melt(id_vars=["timestamp", "id"], value_vars=["x", "y", "speed", "type"], var_name="var")
     df = df.pivot(columns="id", index=["timestamp", "var"], values="value")
     df = df.reset_index()
     df["timestamp"] = df["timestamp"].dt.strftime("%H:%M:%S")