Simplify enriched_datetime

Author: torotil

sqlalchemy_utils/__init__.py

@@ -68,8 +68,6 @@
     DateTimeRangeType,
     EmailType,
     EncryptedType,
-    EnrichedDateTimeType,
-    EnrichedDateType,
     instrumented_list,
     InstrumentedList,
     Int8RangeType,

sqlalchemy_utils/types/__init__.py

@@ -12,7 +12,6 @@
     EncryptedType,
     StringEncryptedType
 )
-from .enriched_datetime.enriched_date_type import EnrichedDateType  # noqa
 from .ip_address import IPAddressType  # noqa
 from .json import JSONType  # noqa
 from .locale import LocaleType  # noqa
@@ -42,8 +41,6 @@
 from .uuid import UUIDType  # noqa
 from .weekdays import WeekDaysType  # noqa
 
-from .enriched_datetime.enriched_datetime_type import EnrichedDateTimeType  # noqa isort:skip
-
 
 class InstrumentedList(_InstrumentedList):
     """Enhanced version of SQLAlchemy InstrumentedList. Provides some

sqlalchemy_utils/types/arrow.py

@@ -1,62 +1 @@
-from ..exceptions import ImproperlyConfigured
-from .enriched_datetime import ArrowDateTime
-from .enriched_datetime.enriched_datetime_type import EnrichedDateTimeType
-
-arrow = None
-try:
-    import arrow
-except ImportError:
-    pass
-
-
-class ArrowType(EnrichedDateTimeType):
-    """
-    ArrowType provides way of saving Arrow_ objects into database. It
-    automatically changes Arrow_ objects to datetime objects on the way in and
-    datetime objects back to Arrow_ objects on the way out (when querying
-    database). ArrowType needs Arrow_ library installed.
-
-    .. _Arrow: https://github.com/arrow-py/arrow
-
-    ::
-
-        from datetime import datetime
-        from sqlalchemy_utils import ArrowType
-        import arrow
-
-
-        class Article(Base):
-            __tablename__ = 'article'
-            id = sa.Column(sa.Integer, primary_key=True)
-            name = sa.Column(sa.Unicode(255))
-            created_at = sa.Column(ArrowType)
-
-
-
-        article = Article(created_at=arrow.utcnow())
-
-
-    As you may expect all the arrow goodies come available:
-
-    ::
-
-
-        article.created_at = article.created_at.replace(hours=-1)
-
-        article.created_at.humanize()
-        # 'an hour ago'
-
-    """
-    cache_ok = True
-
-    def __init__(self, *args, **kwargs):
-        if not arrow:
-            raise ImproperlyConfigured(
-                "'arrow' package is required to use 'ArrowType'"
-            )
-
-        super().__init__(
-            datetime_processor=ArrowDateTime,
-            *args,
-            **kwargs
-        )
+from .enriched_datetime import ArrowDateTime as ArrowType

sqlalchemy_utils/types/enriched_datetime/arrow_datetime.py

@@ -1,6 +1,9 @@
 from collections.abc import Iterable
 from datetime import datetime
 
+from sqlalchemy import types
+
+from ..scalar_coercible import ScalarCoercible
 from ...exceptions import ImproperlyConfigured
 
 arrow = None
@@ -10,14 +13,36 @@
     pass
 
 
-class ArrowDateTime:
-    def __init__(self):
+class ArrowDateTime(types.TypeDecorator, ScalarCoercible):
+    """Datetime column that uses arrow dates.
+
+    Example::
+
+
+        from sqlalchemy_utils.enriched_datetime import ArrowDateTime
+        import arrow
+
+        class User(Base):
+            __tablename__ = 'user'
+            id = sa.Column(sa.Integer, primary_key=True)
+            created_at = sa.Column(ArrowDateTime)
+
+        user = User()
+        user.created_at = arrow.now()
+        session.add(user)
+        session.commit()
+    """
+    impl = types.DateTime
+    cache_ok = True
+
+    def __init__(self, *args, **kwargs):
         if not arrow:
             raise ImproperlyConfigured(
                 "'arrow' package is required to use 'ArrowDateTime'"
             )
+        super().__init__(*args, **kwargs)
 
-    def _coerce(self, impl, value):
+    def _coerce(self, value):
         if isinstance(value, str):
             value = arrow.get(value)
         elif isinstance(value, Iterable):
@@ -26,14 +51,17 @@ def _coerce(self, impl, value):
             value = arrow.get(value)
         return value
 
-    def process_bind_param(self, impl, value, dialect):
+    def process_bind_param(self, value, dialect):
         if value:
-            utc_val = self._coerce(impl, value).to('UTC')
+            utc_val = self._coerce(value).to('UTC')
             return utc_val.datetime\
-                if impl.timezone else utc_val.naive
+                if self.impl.timezone else utc_val.naive
         return value
 
-    def process_result_value(self, impl, value, dialect):
+    def process_result_value(self, value, dialect):
         if value:
             return arrow.get(value)
         return value
+
+    def process_literal_param(self, value, dialect):
+        return value

sqlalchemy_utils/types/enriched_datetime/enriched_date_type.py

@@ -1,3 +1,5 @@
+from sqlalchemy import types
+
 from ...exceptions import ImproperlyConfigured
 from .pendulum_datetime import PendulumDateTime
 
@@ -9,24 +11,22 @@
 
 
 class PendulumDate(PendulumDateTime):
-    def __init__(self):
-        if not pendulum:
-            raise ImproperlyConfigured(
-                "'pendulum' package is required to use 'PendulumDate'"
-            )
 
-    def _coerce(self, impl, value):
+    cache_ok = True
+    impl = types.Date
+
+    def _coerce(self, value):
         if value:
             if not isinstance(value, pendulum.Date):
-                value = super()._coerce(impl, value).date()
+                value = super()._coerce(value).date()
         return value
 
-    def process_result_value(self, impl, value, dialect):
+    def process_result_value(self, value, dialect):
         if value:
             return pendulum.parse(value.isoformat()).date()
         return value
 
-    def process_bind_param(self, impl, value, dialect):
+    def process_bind_param(self, value, dialect):
         if value:
-            return self._coerce(impl, value)
+            return self._coerce(value)
         return value

sqlalchemy_utils/types/enriched_datetime/enriched_datetime_type.py

@@ -1,5 +1,8 @@
 import datetime
 
+from sqlalchemy import types
+
+from ..scalar_coercible import ScalarCoercible
 from ...exceptions import ImproperlyConfigured
 
 pendulum = None
@@ -9,14 +12,38 @@
     pass
 
 
-class PendulumDateTime:
-    def __init__(self):
+class PendulumDateTime(types.TypeDecorator, ScalarCoercible):
+    """Datetime column that uses PendulumDateTime as values.
+
+    Example::
+
+
+        from sqlalchemy_utils.enriched_datetime import PendulumDateTime
+        import pendulum
+
+        class User(Base):
+            __tablename__ = 'user'
+            id = sa.Column(sa.Integer, primary_key=True)
+            created_at = sa.Column(PendulumDateTime)
+
+        user = User()
+        user.created_at = pendulum.now()
+        session.add(user)
+        session.commit()
+    """
+
+    impl = types.DateTime
+    cache_ok = True
+
+    def __init__(self, *args, **kwargs):
         if not pendulum:
             raise ImproperlyConfigured(
                 "'pendulum' package is required to use 'PendulumDateTime'"
             )
+        super().__init__(*args, **kwargs)
+
 
-    def _coerce(self, impl, value):
+    def _coerce(self, value):
         if value is not None:
             if isinstance(value, pendulum.DateTime):
                 pass
@@ -30,14 +57,17 @@ def _coerce(self, impl, value):
                 value = pendulum.parse(value)
         return value
 
-    def process_bind_param(self, impl, value, dialect):
+    def process_bind_param(self, value, dialect):
         if value:
-            return self._coerce(impl, value).in_tz('UTC').naive()
+            return self._coerce(value).in_tz('UTC').naive()
         return value
 
-    def process_result_value(self, impl, value, dialect):
+    def process_result_value(self, value, dialect):
         if value:
             return pendulum.DateTime.instance(
                 value.replace(tzinfo=datetime.timezone.utc)
             )
         return value
+
+    def process_literal_param(self, value, dialect):
+        return value