chore: better train search (#40)

templates/api/trains_search.html

@@ -1,3 +1,3 @@
 {% for train in trains %}
-  <option value="{{ train.number }}">{{ train.type }} {{ train.number }}</option>
+  <option>{{ train.type }} {{ train.number }}</option>
 {% endfor %}

templates/index.html

@@ -42,7 +42,7 @@ <h1>VerspätungsOrakel</h1>
             <div>
               <label for="train">Zugnummer</label>
               <input
-                type="number"
+                type="text"
                 id="train"
                 name="train"
                 list="trains"

verspaetungsorakel/bahn.py

@@ -33,6 +33,8 @@ def get_important_stations():
     stations = list(filter(lambda s: s.Verkehr == "FV", stations))
     stations = list(filter(lambda s: "Hbf" in s.NAME, stations))
 
+    stations = list(filter(lambda s: "Berlin" in s.NAME, stations))
+
     return stations
 
 

verspaetungsorakel/database.py

@@ -17,9 +17,11 @@ def sqlite_wal_mode(db, connection: sqlite3.Connection):
 
 
 class Train(db.Entity):
-    number = PrimaryKey(str)
+    number = Required(str)
     type = Required(str)
 
+    PrimaryKey(number, type)
+
     trips = Set(lambda: Trip)
 
 

verspaetungsorakel/main.py

@@ -56,17 +56,21 @@ def get_stations(station: str, request: Request):
 
 @app.get("/api/v1/trains")
 @limiter.limit("60/minute")
-def get_trains(train: int, request: Request):
+def get_trains(train: str, request: Request):
+    t_type, t_number = upack_train(train)
+
     with db_session:
-        trains = to_dicts(Train.select(lambda t: str(train) in t.number).limit(100)[:])
+        trains = to_dicts(Train.select(lambda t: t_number in t.number and t_type in t.type).limit(100)[:])
     return templates.TemplateResponse("api/trains_search.html", {"request": request, "trains": trains})
 
 
 @app.get("/api/v1/data")
 @limiter.limit("1/second")
-def search(station: str, train: int, request: Request):
+def search(station: str, train: str, request: Request):
+    t_type, t_number = upack_train(train)
+
     with db_session:
-        db_train = Train[str(train)]
+        db_train = Train.get(number=t_number, type=t_type)
         if not db_train:
             raise HTTPException(status_code=404, detail="Train not found")
         db_station = Station.get(name=station)
@@ -97,3 +101,14 @@ def search(station: str, train: int, request: Request):
 
 def to_dicts(entities: list) -> list:
     return [e.to_dict() for e in entities]
+
+
+def upack_train(train: str) -> tuple:
+    t_args = train.split(" ")
+    if len(t_args) != 2:
+        raise HTTPException(status_code=400, detail="Invalid train number")
+
+    t_type = t_args[0].upper()
+    t_number = t_args[1]
+
+    return t_type, t_number