From 1c2f41587b05a4db97ff5330dbaf0be3938dc8fa Mon Sep 17 00:00:00 2001 From: Kumi Date: Fri, 18 Aug 2023 09:38:18 +0200 Subject: [PATCH] chore: Add docstrings to all functions chore: Add missing type hints --- trackbert.py | 125 ++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 115 insertions(+), 10 deletions(-) diff --git a/trackbert.py b/trackbert.py index 41a06bb..8e0fc6a 100644 --- a/trackbert.py +++ b/trackbert.py @@ -7,6 +7,8 @@ import subprocess import argparse import logging +from typing import Tuple, Never + # Print date and time and level with message logging.basicConfig( format="%(asctime)s %(levelname)s: %(message)s", @@ -14,21 +16,43 @@ logging.basicConfig( datefmt="%Y-%m-%d %H:%M:%S", ) -def notify(title, message): +def notify(title: str, message: str): + """Send a desktop notification + + If notify-send is not found, this function will do nothing. + + Args: + title (str): The title of the notification + message (str): The message of the notification + """ + logging.debug(f"Sending notification: {title} - {message}") + try: subprocess.run( [ "notify-send", title, message, + # TODO: Look into other options, like adding an icon ] ) + except FileNotFoundError: + # If notify-send is not found, do nothing logging.warning("notify-send not found, not sending notification") -def create_shipment(db, tracking_number: str, carrier: str, description=""): +def create_shipment(db: sqlite3.Connection, tracking_number: str, carrier: str, description: str = ""): + """Create a shipment in the database + + Args: + db (sqlite3.Connection): The database connection + tracking_number (str): The tracking number + carrier (str): The carrier slug (e.g. "ups", "austrian_post") + description (str, optional): A description for the shipment, displayed in place of the tracking number in notifications. Defaults to "". + """ + logging.debug(f"Creating shipment for {tracking_number} with carrier {carrier}") db.execute( "INSERT INTO shipments (tracking_number, carrier, description) VALUES (?, ?, ?)", @@ -37,21 +61,50 @@ def create_shipment(db, tracking_number: str, carrier: str, description=""): db.commit() -def get_shipment(db: sqlite3.Connection, tracking_number: str): +def get_shipment(db: sqlite3.Connection, tracking_number: str) -> Tuple[int, str, str, str]: + """Get a shipment from the database + + Args: + db (sqlite3.Connection): The database connection + tracking_number (str): The tracking number + + Returns: + Tuple[int, str, str, str]: The shipment (id, tracking_number, carrier, description) + """ + logging.debug(f"Getting shipment for {tracking_number}") cur = db.cursor() cur.execute("SELECT * FROM shipments WHERE tracking_number = ?", (tracking_number,)) return cur.fetchone() -def get_shipments(db: sqlite3.Connection): +def get_shipments(db: sqlite3.Connection) -> Tuple[Tuple[int, str, str, str]]: + """Get all shipments from the database + + Args: + db (sqlite3.Connection): The database connection + + Returns: + Tuple[Tuple[int, str, str, str]]: All shipments (id, tracking_number, carrier, description) + """ + logging.debug(f"Getting all shipments") cur = db.cursor() cur.execute("SELECT * FROM shipments") return cur.fetchall() -def get_shipment_events(db, shipment_id): +def get_shipment_events(db, shipment_id) -> Tuple[Tuple[int, int, str, str, str]]: + """Get all events for a shipment from the database + + Args: + db (sqlite3.Connection): The database connection + shipment_id (int): The shipment id + + Returns: + Tuple[Tuple[int, int, str, str, str]]: All events for the shipment (id, shipment_id, event_time, event_description, raw_event) + """ + logging.debug(f"Getting events for shipment {shipment_id}") cur = db.cursor() cur.execute("SELECT * FROM events WHERE shipment_id = ?", (shipment_id,)) @@ -65,6 +118,16 @@ def create_event( event_description, raw_event, ): + """Create an event for a shipment in the database + + Args: + db (sqlite3.Connection): The database connection + shipment_id (int): The shipment id + event_time (str): The event time + event_description (str): The event description + raw_event (str): The raw event + """ + logging.debug(f"Creating event for shipment {shipment_id}: {event_description} - {event_time}") db.execute( "INSERT INTO events (shipment_id, event_time, event_description, raw_event) VALUES (?, ?, ?, ?)", @@ -78,7 +141,17 @@ def create_event( db.commit() -def get_latest_event(db, shipment_id): +def get_latest_event(db, shipment_id) -> Tuple[int, int, str, str, str]: + """Get the latest event for a shipment from the database + + Args: + db (sqlite3.Connection): The database connection + shipment_id (int): The shipment id + + Returns: + Tuple[int, int, str, str, str]: The latest event (id, shipment_id, event_time, event_description, raw_event) + """ + logging.debug(f"Getting latest event for shipment {shipment_id}") cur = db.cursor() cur.execute( @@ -89,6 +162,12 @@ def get_latest_event(db, shipment_id): def initialize_db(db): + """Initialize the database - create tables if they don't exist + + Args: + db (sqlite3.Connection): The database connection + """ + logging.debug("Initializing database") db.execute( "CREATE TABLE IF NOT EXISTS shipments (id INTEGER PRIMARY KEY AUTOINCREMENT, tracking_number TEXT, carrier TEXT, description TEXT)" @@ -99,14 +178,30 @@ def initialize_db(db): db.commit() -def get_db(): +def get_db(path: str = "trackbert.db"): + """Get a database connection + + Args: + path (str, optional): The path to the database file. Defaults to "trackbert.db" in the current directory. + + Returns: + sqlite3.Connection: The database connection + """ + logging.debug("Connecting to database") - db = sqlite3.connect("trackbert.db") + db = sqlite3.connect(path) initialize_db(db) return db -def start_loop(db, api: KeyDelivery): +def start_loop(db: sqlite3.Connection, api: KeyDelivery) -> Never: + """Start the main loop + + Args: + db (sqlite3.Connection): The database connection + api (KeyDelivery): The KeyDelivery API object + """ + logging.debug("Starting loop") while True: for shipment in get_shipments(db): @@ -142,7 +237,9 @@ def start_loop(db, api: KeyDelivery): time.sleep(300) -def main(): +def main() -> Never: + """Main function - get the database connection, create the KeyDelivery API object, and start the main loop""" + db = get_db() api = KeyDelivery.from_config("config.ini") notify("Trackbert", "Starting up") @@ -150,6 +247,8 @@ def main(): if __name__ == "__main__": + # Parse command line arguments + parser = argparse.ArgumentParser() parser.add_argument("--tracking-number", "-n", type=str, required=False) @@ -157,14 +256,20 @@ if __name__ == "__main__": parser.add_argument("--description", "-d", type=str, required=False) args = parser.parse_args() + # If the user specified a tracking number and carrier, create a shipment and exit + if args.tracking_number is not None and args.carrier is not None: db = get_db() create_shipment(db, args.tracking_number, args.carrier, args.description) print(f"Created shipment for {args.tracking_number} with carrier {args.carrier}") exit(0) + # If the user specified a tracking number but not a carrier, error out + if args.tracking_number is not None: print("You must specify a carrier with -c") exit(1) + # If no arguments were specified, start the main loop + main()