feat: Add write_event method to Database class

This commit adds a new method called `write_event` to the `Database` class in `database.py`. The method takes an `Event` object as a parameter and adds it to the session before committing the changes to the database. This change simplifies the code in the `create_event` method by making use of the new `write_event` method.

The `write_event` method is used in the `Tracker` class in `tracker.py` to replace the call to `create_event`. This change improves the clarity and readability of the code.

Additionally, in the `Tracker` class, the `api` object is no longer created using `KeyDelivery.from_config`, as the `api` object is no longer referenced in the code. This change is made to remove unused code and improve code quality.

This commit also adds some new import statements to the `Tracker` class in `tracker.py` to support the new changes.

Lastly, two new tracker classes `KeyDelivery` and `PostAT` are added under the `trackers` directory. These new classes inherit from the `BaseTracker` class and implement the `get_status` method. The `KeyDelivery` class makes use of the `pykeydelivery.KeyDelivery` API to fetch tracking events, while the `PostAT` class uses the `postat.classes.api.PostAPI` API to fetch tracking events. The new tracker classes also define the supported carriers and their priorities.

Overall, these changes improve the functionality and structure of the codebase.

classes/database.py

@@ -50,7 +50,10 @@ class Database:
             raw_event = json.dumps(raw_event)
 
         new_event = Event(shipment_id=shipment_id, event_time=event_time, event_description=event_description, raw_event=raw_event)
-        self.session.add(new_event)
+        self.write_event(new_event)
+
+    def write_event(self, event):
+        self.session.add(event)
         self.session.commit()
 
     def get_shipment_events(self, shipment_id):

classes/tracker.py

@@ -1,10 +1,12 @@
 import logging
 import subprocess
 import time
+import importlib
 from pathlib import Path
 from typing import Optional, Tuple, Never
 
 from .database import Database
+from trackers.base import BaseTracker
 
 from pykeydelivery import KeyDelivery
 
@@ -16,6 +18,41 @@ class Tracker:
             datefmt="%Y-%m-%d %H:%M:%S",
         )
 
+        self.find_apis()
+
+    def find_apis(self):
+        logging.debug("Finding APIs")
+
+        self.apis = []
+
+        for api in Path(__file__).parent.parent.glob("trackers/*.py"):
+            if api.name in ("__init__.py", "base.py"):
+                continue
+
+            logging.debug(f"Found API {api.stem}")
+
+            module = importlib.import_module(f"trackers.{api.stem}")
+
+            if "tracker" in module.__dict__:
+                tracker = module.tracker
+                logging.debug(f"Found tracker {api.stem}")
+                try:
+                    carriers = tracker.supported_carriers()
+                    api = tracker()
+
+                    for carrier, priority in carriers:
+                        self.apis.append((carrier, priority, api))
+                except:
+                    logging.exception(f"Error loading tracker {name}")
+        
+    def query_api(self, tracking_number: str, carrier: str) -> list:
+        logging.debug(f"Querying API for {tracking_number} with carrier {carrier}")
+
+        for api_carrier, _, api in sorted(self.apis, key=lambda x: x[1], reverse=True):
+            if api_carrier == "*" or api_carrier == carrier:
+                logging.debug(f"Using API {api.__class__.__name__} for {tracking_number} with carrier {carrier}")
+                return list(api.get_status(tracking_number, carrier))
+
     def notify(self, title: str, message: str, urgency: str = "normal", timeout: Optional[int] = 5000) -> None:
         logging.debug(f"Sending notification: {title} - {message}")
 
@@ -39,6 +76,7 @@ class Tracker:
 
     def start_loop(self) -> Never:
         logging.debug("Starting loop")
+
         while True:
             for shipment in self.db.get_shipments():
                 shipment_id = shipment.id
@@ -50,36 +88,23 @@ class Tracker:
 
                 latest_known_event = self.db.get_latest_event(shipment_id)
 
-                all_events = self.api.realtime(carrier, tracking_number)
-
-                try:
-                    logging.debug(f"Got events for {tracking_number}: {len(all_events['data']['items'])}")
-                except KeyError:
-                    print(f"Error getting events for {tracking_number}: {all_events}")
-                    continue
-
-                events = sorted(all_events["data"]["items"], key=lambda x: x["time"], reverse=True)
+                events = self.query_api(tracking_number, carrier)
 
                 if latest_known_event:
                     logging.debug(f"Latest known event for {tracking_number}: {latest_known_event.event_description} - {latest_known_event.event_time}")
                 else:
                     logging.debug(f"No known events for {tracking_number}")
 
-                logging.debug(f"Latest upstream event for {tracking_number}: {events[0]['context']} - {events[0]['time']}")
+                logging.debug(f"Latest upstream event for {tracking_number}: {events[0].event_description} - {events[0].event_time}")
 
                 latest = True
 
                 for event in events:
-                    if latest_known_event is None or event["time"] > latest_known_event.event_time:
-                        self.db.create_event(
-                            shipment_id,
-                            event["time"],
-                            event["context"],
-                            event,
-                        )
+                    if latest_known_event is None or event.event_time > latest_known_event.event_time:
+                        self.db.write_event(event)
 
-                        logging.info(f"New event for {tracking_number}: {event['context']} - {event['time']}")
-                        self.notify(f"New event for {description or tracking_number}", event["context"] + " - " + event["time"], urgency="critical" if latest else "normal")
+                        logging.info(f"New event for {tracking_number}: {event.event_description} - {event.event_time}")
+                        self.notify(f"New event for {description or tracking_number}", event.event_description + " - " + event.event_time, urgency="critical" if latest else "normal")
 
                         latest = False
 
@@ -87,6 +112,5 @@ class Tracker:
 
     def start(self):
         self.db = Database('sqlite:///trackbert.db')
-        self.api = KeyDelivery.from_config("config.ini")
         self.notify("Trackbert", "Starting up")
         self.start_loop()

requirements.txt

@@ -1,3 +1,4 @@
 pykeydelivery
 postat
-sqlalchemy
+sqlalchemy
+python-dateutil

trackers/base.py

@@ -0,0 +1,21 @@
+class BaseTracker:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def get_status(self, tracking_number, carrier):
+        raise NotImplementedError()
+
+    @staticmethod
+    def supported_carriers():
+        """Defines the carriers supported by this tracker.
+
+        Returns:
+            list: List of supported carriers as tuples of (carrier_code, priority), 
+                where priority is an integer. The carrier with the highest priority
+                will be used when tracking a shipment. "*" can be used as a wildcard
+                to match all carriers.
+
+        Raises:
+            NotImplementedError: When this method is not implemented by the subclass.
+        """
+        raise NotImplementedError()

trackers/keydelivery.py

@@ -0,0 +1,39 @@
+from .base import BaseTracker
+from classes.database import Event
+
+from pykeydelivery import KeyDelivery as KeyDeliveryAPI
+
+import json
+import logging
+
+
+class KeyDelivery(BaseTracker):
+    def __init__(self, *args, **kwargs):
+        self.api = KeyDeliveryAPI.from_config("config.ini")
+
+    def get_status(self, tracking_number, carrier):
+        all_events = self.api.realtime(carrier, tracking_number)
+
+        try:
+            logging.debug(f"Got events for {tracking_number}: {len(all_events['data']['items'])}")
+        except KeyError:
+            print(f"Error getting events for {tracking_number}: {all_events}")
+            return
+
+        events = sorted(all_events["data"]["items"], key=lambda x: x["time"], reverse=True)
+
+        for event in events:
+            yield Event(
+                shipment_id = tracking_number,
+                event_time = event["time"],
+                event_description = event["context"],
+                raw_event = json.dumps(event)
+            )
+        
+    @staticmethod
+    def supported_carriers():
+        return [
+            ("*", 1),
+        ]
+
+tracker = KeyDelivery

trackers/postat.py

@@ -0,0 +1,36 @@
+from .base import BaseTracker
+from classes.database import Event
+
+import json
+
+from dateutil.parser import parse
+from postat.classes.api import PostAPI
+
+class PostAT(BaseTracker):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def get_status(self, tracking_number, carrier):
+        api = PostAPI()
+        status = api.get_shipment_status(tracking_number)
+        shipment = status["data"]["einzelsendung"]
+        events = shipment["sendungsEvents"]
+
+        for event in events:
+            timestamp = event["timestamp"]
+            py_timestamp = parse(timestamp)
+            event_time = py_timestamp.strftime("%Y-%m-%d %H:%M:%S")
+            yield Event(
+                shipment_id = tracking_number,
+                event_time = event_time,
+                event_description = event["text"],
+                raw_event = json.dumps(event)
+            )
+
+    @staticmethod
+    def supported_carriers():
+        return [
+            ("austrian_post", 100),
+        ]
+
+tracker = PostAT