chore: Add docstrings to all functions

chore: Add missing type hints
This commit is contained in:
Kumi 2023-08-18 09:38:18 +02:00
parent 018cccb6a1
commit 1c2f41587b
Signed by: kumi
GPG key ID: ECBCC9082395383F

View file

@ -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()