From bd8f5280112b2cf1cf6113453977629be2ee92c5 Mon Sep 17 00:00:00 2001
From: Kevin Morris <kevr@0cost.org>
Date: Sun, 13 Jun 2021 10:48:31 -0700
Subject: [PATCH] add Base.as_dict() and Base.json()

Two utility functions for all of our ORM models that will
allow us to easily convert them to Python structures and
JSON data.

Signed-off-by: Kevin Morris <kevr@0cost.org>
---
 aurweb/models/declarative.py | 30 ++++++++++++++++++++++++++++++
 aurweb/util.py               |  8 ++++++++
 test/test_user.py            | 19 +++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/aurweb/models/declarative.py b/aurweb/models/declarative.py
index 45a629ce..96ee1829 100644
--- a/aurweb/models/declarative.py
+++ b/aurweb/models/declarative.py
@@ -1,10 +1,40 @@
+import json
+
 from sqlalchemy.ext.declarative import declarative_base
 
 import aurweb.db
 
+from aurweb import util
+
+
+def to_dict(model):
+    return {
+        c.name: getattr(model, c.name)
+        for c in model.__table__.columns
+    }
+
+
+def to_json(model, indent: int = None):
+    return json.dumps({
+        k: util.jsonify(v)
+        for k, v in to_dict(model).items()
+    }, indent=indent)
+
+
 Base = declarative_base()
+
+# Setup __table_args__ applicable to every table.
 Base.__table_args__ = {
     "autoload": True,
     "autoload_with": aurweb.db.get_engine(),
     "extend_existing": True
 }
+
+
+# Setup Base.as_dict and Base.json.
+#
+# With this, declarative models can use .as_dict() or .json()
+# at any time to produce a dict and json out of table columns.
+#
+Base.as_dict = to_dict
+Base.json = to_json
diff --git a/aurweb/util.py b/aurweb/util.py
index 8e4b291d..ad8ac6b7 100644
--- a/aurweb/util.py
+++ b/aurweb/util.py
@@ -3,6 +3,7 @@ import random
 import re
 import string
 
+from datetime import datetime
 from urllib.parse import urlparse
 
 from email_validator import EmailNotValidError, EmailUndeliverableError, validate_email
@@ -94,3 +95,10 @@ def account_url(context, user):
     if request.url.scheme == "http" and request.url.port != 80:
         base += f":{request.url.port}"
     return f"{base}/account/{user.Username}"
+
+
+def jsonify(obj):
+    """ Perform a conversion on obj if it's needed. """
+    if isinstance(obj, datetime):
+        obj = int(obj.timestamp())
+    return obj
diff --git a/test/test_user.py b/test/test_user.py
index 8b4da61e..06585207 100644
--- a/test/test_user.py
+++ b/test/test_user.py
@@ -1,4 +1,5 @@
 import hashlib
+import json
 
 from datetime import datetime, timedelta
 
@@ -198,3 +199,21 @@ def test_user_credential_types():
     assert aurweb.auth.trusted_user(user)
     assert aurweb.auth.developer(user)
     assert aurweb.auth.trusted_user_or_dev(user)
+
+
+def test_user_json():
+    data = json.loads(user.json())
+    assert data.get("ID") == user.ID
+    assert data.get("Username") == user.Username
+    assert data.get("Email") == user.Email
+    # .json() converts datetime values to integer timestamps.
+    assert isinstance(data.get("RegistrationTS"), int)
+
+
+def test_user_as_dict():
+    data = user.as_dict()
+    assert data.get("ID") == user.ID
+    assert data.get("Username") == user.Username
+    assert data.get("Email") == user.Email
+    # .as_dict() does not convert values to json-capable types.
+    assert isinstance(data.get("RegistrationTS"), datetime)