From 5e0d7b174cb48b56aaf64ccabc7097c88c4472ba Mon Sep 17 00:00:00 2001
From: Adrien <me@adrien.run>
Date: Thu, 11 May 2023 21:40:38 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=8F=B7=EF=B8=8F=20Fix=20some=20type=20iss?=
 =?UTF-8?q?ues=20(mypy)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/backend/schemas/next_passage.py |  4 ++--
 backend/routers/stop.py                 | 30 +++++++++++++++++--------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/backend/backend/schemas/next_passage.py b/backend/backend/schemas/next_passage.py
index 68d0298..e895190 100644
--- a/backend/backend/schemas/next_passage.py
+++ b/backend/backend/schemas/next_passage.py
@@ -4,7 +4,7 @@ from ..idfm_interface.idfm_types import TrainStatus
 
 
 class NextPassage(BaseModel):
-    line: str
+    line: int
     operator: str
     destinations: list[str]
     atStop: bool
@@ -19,4 +19,4 @@ class NextPassage(BaseModel):
 
 class NextPassages(BaseModel):
     ts: int
-    passages: dict[str, dict[str, list[NextPassage]]]
+    passages: dict[int, dict[str, list[NextPassage]]]
diff --git a/backend/routers/stop.py b/backend/routers/stop.py
index cca2260..1f3e803 100644
--- a/backend/routers/stop.py
+++ b/backend/routers/stop.py
@@ -4,7 +4,11 @@ from typing import Sequence
 
 from fastapi import APIRouter, HTTPException
 
-from backend.idfm_interface import Destinations as IdfmDestinations, IdfmInterface
+from backend.idfm_interface import (
+    Destinations as IdfmDestinations,
+    IdfmInterface,
+    TrainStatus,
+)
 from backend.models import Stop, StopArea, StopShape
 from backend.schemas import (
     NextPassage as NextPassageSchema,
@@ -38,16 +42,20 @@ def optional_datetime_to_ts(dt: datetime | None) -> int | None:
 @router.get("/")
 async def get_stop(
     name: str = "", limit: int = 10
-) -> Sequence[StopAreaSchema | StopSchema]:
+) -> Sequence[StopAreaSchema | StopSchema] | None:
+
+    matching_stops = await Stop.get_by_name(name)
+    if matching_stops is None:
+        return None
 
     formatted: list[StopAreaSchema | StopSchema] = []
-    matching_stops = await Stop.get_by_name(name)
-
     stop_areas: dict[int, StopArea] = {}
     stops: dict[int, Stop] = {}
     for stop in matching_stops:
-        dst = stop_areas if isinstance(stop, StopArea) else stops
-        dst[stop.id] = stop
+        if isinstance(stop, StopArea):
+            stop_areas[stop.id] = stop
+        elif isinstance(stop, Stop):
+            stops[stop.id] = stop
 
     for stop_area in stop_areas.values():
 
@@ -121,8 +129,12 @@ async def get_next_passages(stop_id: int) -> NextPassagesSchema | None:
                 arrivalPlatformName=arrivalPlatformName,
                 aimedDepartTs=optional_datetime_to_ts(call.AimedDepartureTime),
                 expectedDepartTs=optional_datetime_to_ts(call.ExpectedDepartureTime),
-                arrivalStatus=call.ArrivalStatus.value,
-                departStatus=call.DepartureStatus.value,
+                arrivalStatus=call.ArrivalStatus
+                if call.ArrivalStatus is not None
+                else TrainStatus.unknown,
+                departStatus=call.DepartureStatus
+                if call.DepartureStatus is not None
+                else TrainStatus.unknown,
             )
 
             by_line_passages = by_line_by_dst_passages[line_id]
@@ -131,7 +143,7 @@ async def get_next_passages(stop_id: int) -> NextPassagesSchema | None:
                 by_line_passages[dst].append(next_passage)
 
     return NextPassagesSchema(
-        ts=service_delivery.ResponseTimestamp.timestamp(),
+        ts=int(service_delivery.ResponseTimestamp.timestamp()),
         passages=by_line_by_dst_passages,
     )