Umbrellas.read can optionally return dict for single umbrella

1 files changed, 24 insertions, 22 deletions

diff --git a/jimbrella/umbrellas.py b/jimbrella/umbrellas.py
index 6da7184..ff88277 100644
--- a/jimbrella/umbrellas.py
+++ b/jimbrella/umbrellas.py
@@ -1,5 +1,6 @@
 import sqlite3
 from datetime import datetime, timedelta
+from typing import Union
 from .utils import human_datetime, human_timedelta
 from .config import DUE_HOURS, ADMIN_LOG_PATH
 from .exceptions import *
@@ -42,12 +43,21 @@ class Umbrellas:
         """
         self.path = path
 
-    def read(self) -> list:
+    def read(self, umbid=None) -> Union[dict, list]:
+        """Read umbrella data from database.
+
+        If umbid is an integer, returns dict pertaining to umbrella #<umbid>.
+        If umbid is None, returns list of dicts for all umbrellas.
+        """
         db = sqlite3.connect(self.path)
-        db.row_factory = sqlite.Row
-        umbrellas = db.execute("SELECT * FROM Umbrellas").fetchall()
+        db.row_factory = sqlite3.Row
+        if umbid is None:
+            data = db.execute("SELECT * FROM Umbrellas").fetchall()
+        else:
+            data = db.execute("SELECT * FROM Umbrellas WHERE id = ?", (umbid,)).fetchone()
+
         db.close()
-        return umbrellas
+        return data
 
     def update(self, umb) -> dict:
         """Update Umbrella table with new data given in `umb`.
@@ -76,7 +86,7 @@ class Umbrellas:
 
         # check if umbrella #<id> exists in database
         umbid = umb["id"]
-        umb_in_db = db.execute("SELECT * FROM Umbrellas WHERE id = ?", umbid).fetchone()
+        umb_in_db = db.execute("SELECT * FROM Umbrellas WHERE id = ?", (umbid,)).fetchone()
         if umb_in_db is None:
             raise UmbrellaNotFoundError(umbid)
 
@@ -89,7 +99,7 @@ class Umbrellas:
                 if umb_in_db["status"] != umb["status"]:
                     diff["status"] = (umb_in_db["status"], umb["status"])
 
-                db.execute("UPDATE Umbrellas SET status = ? WHERE id = ?", status, umbid)
+                db.execute("UPDATE Umbrellas SET status = ? WHERE id = ?", (status, umbid))
             else:
                 raise UmbrellaValueError("status")
 
@@ -106,9 +116,9 @@ class Umbrellas:
 
                     db.execute(
                         "UPDATE Umbrellas SET ? = ? WHERE id = ?",
-                        col,
+                        (col,
                         umb[col],
-                        umbid,
+                        umbid,)
                     )
 
             if "lent_at" in umb:
@@ -130,8 +140,8 @@ class Umbrellas:
 
                 db.execute(
                     "UPDATE Umbrellas SET lent_at = ? WHERE id = ?",
-                    lent_at.isoformat(timespec="milliseconds"),
-                    umbid,
+                   ( lent_at.isoformat(timespec="milliseconds"),
+                    umbid,)
                 )
         else:
             # discard unneeded fields
@@ -142,7 +152,7 @@ class Umbrellas:
                 "tenant_email",
                 "lent_at",
             ):
-                db.execute("UPDATE Umbrellas SET ? = NULL WHERE id = ?", col, umbid)
+                db.execute("UPDATE Umbrellas SET ? = NULL WHERE id = ?", (col, umbid))
 
         # now that new data are validated, commit the SQL transaction
         db.commit()
@@ -153,10 +163,7 @@ class Umbrellas:
         self, umbid, date, tenant_name, tenant_id, tenant_phone="", tenant_email=""
     ) -> None:
         """When a user has borrowed an umbrella."""
-        db = sqlite3.connect(self.path)
-        db.row_factory = sqlite3.Row
-        umb = db.execute("SELECT * FROM Umbrellas WHERE id = ?", umbid)
-        db.close()
+        umb = self.read(umbid)
 
         if umb is None:
             raise UmbrellaNotFoundError(umbid)
@@ -181,10 +188,7 @@ class Umbrellas:
         `tenant_name` and `tenant_id` are used to verify if the umbrella is returned by the same
         person who borrowed it.
         """
-        db = sqlite3.connect(self.path)
-        db.row_factory = sqlite3.Row
-        umb = db.execute("SELECT * FROM Umbrellas WHERE id = ?", umbid)
-        db.close()
+        umb = self.read(umbid)
 
         if umb is None:
             raise UmbrellaNotFoundError(umbid)
@@ -202,9 +206,7 @@ class Umbrellas:
 
     def mark_overdue(self, umbid) -> None:
         """When an umbrella is overdue, change its status to "overdue"."""
-        db = sqlite3.connect(self.path)
-        db.row_factory = sqlite3.Row
-        umb = db.execute("SELECT * FROM Umbrellas WHERE id = ?", umbid)
+        umb = self.read(umbid)
 
         if umb is None:
             raise UmbrellaNotFoundError(umbid)