From ab79366ec5220475b0dc4f3440a81caca836de18 Mon Sep 17 00:00:00 2001
From: "0x90.space" <people@schleuder.0x90.space>
Date: Wed, 29 Dec 2021 02:52:08 +0100
Subject: [PATCH] provide api

---
 setup.cfg                      |  1 +
 src/ticketfrei3/application.py | 74 ++++++++++++++++++++++++++++++----
 src/ticketfrei3/models.py      | 12 +++++-
 3 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index 61d67a7..4731e15 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -22,6 +22,7 @@ python_requires = >=3.8
 install_requires =
     aiosqlite
     blacksheep
+    pydantic
     tortoise-orm
     uvicorn
 
diff --git a/src/ticketfrei3/application.py b/src/ticketfrei3/application.py
index 58d2a17..4b14853 100644
--- a/src/ticketfrei3/application.py
+++ b/src/ticketfrei3/application.py
@@ -1,9 +1,23 @@
 from blacksheep.client import ClientSession
+from blacksheep.contents import Content
+from blacksheep.messages import Response
 from blacksheep.server import Application
-import json
-from ticketfrei3.models import Fahrt, Halt, Haltestelle
+from blacksheep.server.responses import not_found
+from json import loads as from_json
+from ticketfrei3.models import (
+    Fahrt,
+    Fahrt_Pydantic,
+    Fahrt_Pydantic_List,
+    Halt,
+    Halt_Pydantic,
+    Halt_Pydantic_List,
+    Haltestelle,
+    Haltestelle_Pydantic,
+    Haltestelle_Pydantic_List,
+)
 from tortoise import Tortoise
 from tortoise.transactions import in_transaction
+from tortoise.contrib.pydantic import PydanticModel
 
 
 app = Application()
@@ -29,10 +43,10 @@ async def finalize(app: Application) -> None:
 app.on_stop += finalize
 
 
-@app.route("/fetch_haltestellen")
+@app.route("/api/fetch_haltestellen")
 async def fetch_haltestellen(client: ClientSession) -> None:
     response = await client.get("https://start.vag.de/dm/api/v1/haltestellen.json/vag")
-    haltestellen = json.loads(await response.text())
+    haltestellen = from_json(await response.text())
     for haltestelle in haltestellen["Haltestellen"]:
         if await Haltestelle.exists(id=haltestelle["VGNKennung"]):
             continue
@@ -44,13 +58,13 @@ async def fetch_haltestellen(client: ClientSession) -> None:
         )
 
 
-@app.route("/fetch_fahrten")
+@app.route("/api/fetch_fahrten")
 async def fetch_fahrten(client: ClientSession) -> None:
     for zweig in ("bus", "tram", "ubahn"):
         response = await client.get(
             "https://start.vag.de/dm/api/v1/fahrten.json/%s?timespan=120" % zweig
         )
-        fahrten = json.loads(await response.text())
+        fahrten = from_json(await response.text())
         for fahrt in fahrten["Fahrten"]:
             if await Fahrt.exists(id=fahrt["Fahrtnummer"]):
                 continue
@@ -61,7 +75,7 @@ async def fetch_fahrten(client: ClientSession) -> None:
                     "https://start.vag.de/dm/api/v1/fahrten.json/%s/%d"
                     % (zweig, fahrt.id)
                 )
-                details = json.loads(await response.text())
+                details = from_json(await response.text())
                 for halt in details["Fahrtverlauf"]:
                     haltestelle = (
                         await Haltestelle.filter(id=halt["VGNKennung"])
@@ -74,3 +88,49 @@ async def fetch_fahrten(client: ClientSession) -> None:
                         ankunft=halt.get("AnkunftszeitIst"),
                         abfahrt=halt.get("AbfahrtszeitIst"),
                     ).save(using_db=connection)
+
+
+def to_json(model: PydanticModel) -> Response:
+    return Response(
+        200,
+        content=Content(b"application/json", model.json(ensure_ascii=False).encode()),
+    )
+
+
+@app.router.get("/api/haltestellen")
+async def haltestellen() -> Response:
+    return to_json(await Haltestelle_Pydantic_List.from_queryset(Haltestelle.all()))
+
+
+@app.router.get("/api/haltestelle/{id}")
+async def haltestelle(id: int) -> Response:
+    haltestelle = await Haltestelle.filter(id=id).first()
+    if haltestelle is None:
+        return not_found()
+    return to_json(await Haltestelle_Pydantic.from_tortoise_orm(haltestelle))
+
+
+@app.router.get("/api/fahrten")
+async def fahrten() -> Response:
+    return to_json(await Fahrt_Pydantic_List.from_queryset(Fahrt.all()))
+
+
+@app.router.get("/api/fahrt/{id}")
+async def fahrt(id: int) -> Response:
+    fahrt = await Fahrt.filter(id=id).first()
+    if fahrt is None:
+        return not_found()
+    return to_json(await Fahrt_Pydantic.from_tortoise_orm(fahrt))
+
+
+@app.router.get("/api/halte")
+async def halte() -> Response:
+    return to_json(await Halt_Pydantic_List.from_queryset(Halt.all()))
+
+
+@app.router.get("/api/halt/{id}")
+async def halt(id: int) -> Response:
+    halt = await Halt.filter(id=id).first()
+    if halt is None:
+        return not_found()
+    return to_json(await Halt_Pydantic.from_tortoise_orm(halt))
diff --git a/src/ticketfrei3/models.py b/src/ticketfrei3/models.py
index bd56004..ebc1a96 100644
--- a/src/ticketfrei3/models.py
+++ b/src/ticketfrei3/models.py
@@ -1,5 +1,6 @@
-from tortoise.models import Model
 from tortoise import fields
+from tortoise.contrib.pydantic import pydantic_model_creator, pydantic_queryset_creator
+from tortoise.models import Model
 
 
 class Fahrt(Model):
@@ -22,3 +23,12 @@ class Haltestelle(Model):
     longitude = fields.FloatField(null=True)
     latitude = fields.FloatField(null=True)
     fahrten: fields.ReverseRelation["Halt"]
+
+
+Fahrt_Pydantic = pydantic_model_creator(Fahrt)
+Halt_Pydantic = pydantic_model_creator(Halt)
+Haltestelle_Pydantic = pydantic_model_creator(Haltestelle)
+
+Fahrt_Pydantic_List = pydantic_queryset_creator(Fahrt)
+Halt_Pydantic_List = pydantic_queryset_creator(Halt)
+Haltestelle_Pydantic_List = pydantic_queryset_creator(Haltestelle)