db type that supports aware datetimes in user-land

2 files changed, 36 insertions, 0 deletions

diff --git a/metricity/utils.py b/metricity/utils.py
new file mode 100644
index 0000000..3f3547c
--- /dev/null
+++ b/metricity/utils.py
@@ -0,0 +1,35 @@
+"""General utility functions and classes for Metricity."""
+from datetime import datetime, timezone
+
+from sqlalchemy.engine import Dialect
+from sqlalchemy.types import DateTime, TypeDecorator
+
+
+class TZDateTime(TypeDecorator):
+    """
+    A db type that supports the use of aware datetimes in user-land.
+
+    Source from SQLAlchemy docs:
+    https://docs.sqlalchemy.org/en/14/core/custom_types.html#store-timezone-aware-timestamps-as-timezone-naive-utc
+
+    Editted to include docstrings and type hints.
+    """
+
+    impl = DateTime
+    cache_ok = True
+
+    def process_bind_param(self, value: datetime, dialect: Dialect) -> datetime:
+        """Convert the value to aware before saving to db."""
+        if value is not None:
+            if not value.tzinfo:
+                raise TypeError("tzinfo is required")
+            value = value.astimezone(timezone.utc).replace(
+                tzinfo=None
+            )
+        return value
+
+    def process_result_value(self, value: datetime, dialect: Dialect) -> datetime:
+        """Convert the value to aware before passing back to user-land."""
+        if value is not None:
+            value = value.replace(tzinfo=timezone.utc)
+        return value
diff --git a/tox.ini b/tox.ini
index 2bbe42d..7866344 100644
--- a/tox.ini
+++ b/tox.ini
@@ -3,3 +3,4 @@ max-line-length=120
 application-import-names=metricity
 import-order-style=pycharm
 exclude=alembic
+ignore=ANN101  # self params in classes.