From c09d8a64f46f202c82b24e41ad5855f4040e938a Mon Sep 17 00:00:00 2001
From: "0x90.space" <people@schleuder.0x90.space>
Date: Wed, 29 Dec 2021 01:01:13 +0100
Subject: [PATCH] refactor and fetch trips and stop times

---
 setup.cfg                      |  1 +
 src/ticketfrei3/__init__.py    | 25 +----------
 src/ticketfrei3/application.py | 76 ++++++++++++++++++++++++++++++++++
 src/ticketfrei3/models.py      | 17 ++++++++
 src/uvicorn/__init__.pyi       |  1 +
 5 files changed, 97 insertions(+), 23 deletions(-)
 create mode 100644 src/ticketfrei3/application.py
 create mode 100644 src/uvicorn/__init__.pyi

diff --git a/setup.cfg b/setup.cfg
index 8b13767..61d67a7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -23,6 +23,7 @@ install_requires =
     aiosqlite
     blacksheep
     tortoise-orm
+    uvicorn
 
 [options.packages.find]
 where = src
diff --git a/src/ticketfrei3/__init__.py b/src/ticketfrei3/__init__.py
index 3938631..5ca4be7 100644
--- a/src/ticketfrei3/__init__.py
+++ b/src/ticketfrei3/__init__.py
@@ -1,26 +1,5 @@
-import asyncio
-from blacksheep.client import ClientSession
-import json
-from ticketfrei3.models import Haltestelle
-from tortoise import Tortoise
-
-
-async def asyncmain() -> None:
-    await Tortoise.init(
-        db_url="sqlite://db.sqlite3", modules={"models": ["ticketfrei3.models"]}
-    )
-    await Tortoise.generate_schemas()
-    async with ClientSession(connection_timeout=30) as client:
-        response = await client.get("https://start.vag.de/dm/api/haltestellen.json/vag")
-        haltestellen = json.loads(await response.text())
-        for haltestelle in haltestellen["Haltestellen"]:
-            if not await Haltestelle.exists(id=haltestelle["VGNKennung"]):
-                await Haltestelle.create(
-                    id=haltestelle["VGNKennung"],
-                    name=haltestelle["Haltestellenname"],
-                )
-    await Tortoise.close_connections()
+import uvicorn
 
 
 def main() -> None:
-    asyncio.run(asyncmain(), debug=True)
+    uvicorn.run("ticketfrei3.application:app")
diff --git a/src/ticketfrei3/application.py b/src/ticketfrei3/application.py
new file mode 100644
index 0000000..58d2a17
--- /dev/null
+++ b/src/ticketfrei3/application.py
@@ -0,0 +1,76 @@
+from blacksheep.client import ClientSession
+from blacksheep.server import Application
+import json
+from ticketfrei3.models import Fahrt, Halt, Haltestelle
+from tortoise import Tortoise
+from tortoise.transactions import in_transaction
+
+
+app = Application()
+
+
+async def init(app: Application) -> None:
+    await Tortoise.init(
+        db_url="sqlite://db.sqlite3", modules={"models": ["ticketfrei3.models"]}
+    )
+    await Tortoise.generate_schemas()
+    app.services.add_instance(ClientSession(connection_timeout=30))
+
+
+app.on_start += init
+
+
+async def finalize(app: Application) -> None:
+    client = app.service_provider.get(ClientSession)
+    await client.close()
+    await Tortoise.close_connections()
+
+
+app.on_stop += finalize
+
+
+@app.route("/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())
+    for haltestelle in haltestellen["Haltestellen"]:
+        if await Haltestelle.exists(id=haltestelle["VGNKennung"]):
+            continue
+        await Haltestelle.create(
+            id=haltestelle["VGNKennung"],
+            name=haltestelle["Haltestellenname"],
+            longitude=haltestelle.get("Longitude"),
+            latitude=haltestelle.get("Latitude"),
+        )
+
+
+@app.route("/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())
+        for fahrt in fahrten["Fahrten"]:
+            if await Fahrt.exists(id=fahrt["Fahrtnummer"]):
+                continue
+            async with in_transaction() as connection:
+                fahrt = Fahrt(id=fahrt["Fahrtnummer"], linie=fahrt["Linienname"])
+                await fahrt.save(using_db=connection)
+                response = await client.get(
+                    "https://start.vag.de/dm/api/v1/fahrten.json/%s/%d"
+                    % (zweig, fahrt.id)
+                )
+                details = json.loads(await response.text())
+                for halt in details["Fahrtverlauf"]:
+                    haltestelle = (
+                        await Haltestelle.filter(id=halt["VGNKennung"])
+                        .using_db(connection)
+                        .first()
+                    )
+                    await Halt(
+                        fahrt=fahrt,
+                        haltestelle=haltestelle,
+                        ankunft=halt.get("AnkunftszeitIst"),
+                        abfahrt=halt.get("AbfahrtszeitIst"),
+                    ).save(using_db=connection)
diff --git a/src/ticketfrei3/models.py b/src/ticketfrei3/models.py
index 22faeb2..bd56004 100644
--- a/src/ticketfrei3/models.py
+++ b/src/ticketfrei3/models.py
@@ -2,6 +2,23 @@ from tortoise.models import Model
 from tortoise import fields
 
 
+class Fahrt(Model):
+    id = fields.IntField(pk=True)
+    linie = fields.TextField()
+    verlauf: fields.ReverseRelation["Halt"]
+
+
+class Halt(Model):
+    id = fields.IntField(pk=True)
+    fahrt = fields.ForeignKeyField("models.Fahrt", "verlauf")
+    haltestelle = fields.ForeignKeyField("models.Haltestelle", "fahrten")
+    ankunft = fields.DatetimeField(null=True)
+    abfahrt = fields.DatetimeField(null=True)
+
+
 class Haltestelle(Model):
     id = fields.IntField(pk=True)
     name = fields.TextField()
+    longitude = fields.FloatField(null=True)
+    latitude = fields.FloatField(null=True)
+    fahrten: fields.ReverseRelation["Halt"]
diff --git a/src/uvicorn/__init__.pyi b/src/uvicorn/__init__.pyi
new file mode 100644
index 0000000..eb7126a
--- /dev/null
+++ b/src/uvicorn/__init__.pyi
@@ -0,0 +1 @@
+def run(app: str) -> None: ...