Allow datetime.datetime type for Umbrellas.update

1 files changed, 23 insertions, 13 deletions

diff --git a/jimbrella/umbrellas.py b/jimbrella/umbrellas.py
index 05b103c..e59bb54 100644
--- a/jimbrella/umbrellas.py
+++ b/jimbrella/umbrellas.py
@@ -59,6 +59,8 @@ class Umbrellas:
         Invalid values are rejected as an UmbrellaValueError.
 
         If `status` is not "lent" or "overdue", `tenant_*` and `lent_at` are automatically erased.
+
+        `lent_at` may be either an ISO 8601 string or a datetime.datetime object. Must be UTC+8.
         """
         # `id` must be specified.
         try:
@@ -83,24 +85,11 @@ class Umbrellas:
             raise UmbrellaValueError("status")
 
         if status in ("lent", "overdue"):
-            try:
-                # timezone must be +08:00
-                if (
-                    datetime.fromisoformat(umb["lent_at"])
-                    .tzinfo.utcoffset(None)
-                    .seconds
-                    != 28800
-                ):
-                    raise UmbrellaValueError("lent_at")
-            except ValueError:
-                raise UmbrellaValueError("lent_at")
-
             for key in (
                 "tenant_name",
                 "tenant_id",
                 "tenant_phone",
                 "tenant_email",
-                "lent_at",
             ):
                 if col in umb:
                     db.execute(
@@ -109,6 +98,27 @@ class Umbrellas:
                         umb[col] or None,
                         umbid,
                     )
+
+            if "lent_at" in umb:
+                try:
+                    # lent_at could be a string, in which case it is parsed
+                    lent_at = datetime.fromisoformat(umb["lent_at"])
+                except TypeError:
+                    # or it could be a datetime.datetime
+                    lent_at = umb["lent_at"]
+                except ValueError:
+                    # anything else is invalid
+                    raise UmbrellaValueError("lent_at")
+
+                if lent_at.tzinfo.utcoffset(None).seconds != 28800:
+                    # timezone must be +08:00
+                    raise UmbrellaValueError("lent_at")
+
+                db.execute(
+                    "UPDATE Umbrellas SET lent_at = ? WHERE id = ?",
+                    lent_at.isoformat(timespec="milliseconds"),
+                    umbid,
+                )
         else:
             # discard unneeded fields
             for col in (