Use native isoformat/fromisoformat (#141)

* Use native __isoformat/fromisoformat * Fixes * Fix fast-path * Add changelod entry * More tests * Backport "%Y-%m-%d" datetime iso format support
anna-money · Dec 11, 2023 · e1dab98 · e1dab98
1 parent 6868ce4
commit e1dab98
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,8 @@
+## unreleased
+
+* [Use native isoformat/fromisoformat](https://github.com/anna-money/marshmallow-recipe/pull/141)
+
+
 ## v0.0.37(2023-12-11)
 
 * [Cache get_pre_loads results](https://github.com/anna-money/marshmallow-recipe/pull/140)

diff --git a/marshmallow_recipe/fields.py b/marshmallow_recipe/fields.py
@@ -3,6 +3,7 @@
 import dataclasses
 import datetime
 import enum
+import sys
 from typing import Any
 
 import marshmallow as m
@@ -633,10 +634,48 @@ def _deserialize(self, value: Any, attr: Any, data: Any, **kwargs: Any) -> Any:
     StrField = StrFieldV3
 
     class DateTimeFieldV3(m.fields.DateTime):
+        SERIALIZATION_FUNCS = {
+            **m.fields.DateTime.SERIALIZATION_FUNCS,  # type: ignore
+            **(
+                {
+                    "iso": datetime.datetime.isoformat,
+                    "iso8601": datetime.datetime.isoformat,
+                }
+                if sys.version_info >= (3, 11)
+                else {}
+            ),
+        }
+
+        @staticmethod
+        def __fromisoformat_py310_fix(
+            format: str, value: str
+        ) -> datetime.datetime:  # pyright: ignore[reportUnusedFunction]
+            try:
+                return m.fields.DateTime.DESERIALIZATION_FUNCS[format](value)  # type: ignore
+            except ValueError:
+                return datetime.datetime.strptime(value, "%Y-%m-%d")
+
+        DESERIALIZATION_FUNCS = {
+            **m.fields.DateTime.DESERIALIZATION_FUNCS,  # type: ignore
+            **(
+                {
+                    "iso": datetime.datetime.fromisoformat,
+                    "iso8601": datetime.datetime.fromisoformat,
+                }
+                if sys.version_info >= (3, 11)
+                else {
+                    "iso": lambda x: DateTimeFieldV3.__fromisoformat_py310_fix("iso", x),
+                    "iso8601": lambda x: DateTimeFieldV3.__fromisoformat_py310_fix("iso8601", x),
+                }
+            ),
+        }
+
         def _deserialize(self, value: Any, attr: Any, data: Any, **kwargs: Any) -> Any:
             result = super()._deserialize(value, attr, data, **kwargs)
             if result.tzinfo is None:
                 return result.replace(tzinfo=datetime.timezone.utc)
+            if result.tzinfo == datetime.timezone.utc:
+                return result
             return result.astimezone(datetime.timezone.utc)
 
         def _serialize(self, value: Any, attr: Any, obj: Any, **kwargs: Any) -> Any:
@@ -834,10 +873,63 @@ def _deserialize(self, value: Any, attr: Any, data: Any, **_: Any) -> Any:
     StrField = StrFieldV2
 
     class DateTimeFieldV2(m.fields.DateTime):
+        @staticmethod
+        def __isoformat(v: datetime.datetime, *, localtime: bool = False, **kwargs: Any) -> str:
+            assert not localtime
+            assert not kwargs
+            return datetime.datetime.isoformat(v)
+
+        DATEFORMAT_SERIALIZATION_FUNCS = {
+            **m.fields.DateTime.DATEFORMAT_SERIALIZATION_FUNCS,  # type: ignore
+            **(
+                {
+                    "iso": __isoformat,
+                    "iso8601": __isoformat,
+                }
+                if sys.version_info >= (3, 11)
+                else {}
+            ),
+        }
+
+        @staticmethod
+        def __fromisoformat_py310_fix(
+            format: str, value: str
+        ) -> datetime.datetime:  # pyright: ignore[reportUnusedFunction]
+            try:
+                return m.fields.DateTime.DATEFORMAT_DESERIALIZATION_FUNCS[format](value)  # type: ignore
+            except ValueError:
+                return datetime.datetime.strptime(value, "%Y-%m-%d")
+
+        DATEFORMAT_DESERIALIZATION_FUNCS = {
+            **m.fields.DateTime.DATEFORMAT_DESERIALIZATION_FUNCS,  # type: ignore
+            **(
+                {
+                    "iso": datetime.datetime.fromisoformat,
+                    "iso8601": datetime.datetime.fromisoformat,
+                }
+                if sys.version_info >= (3, 11)
+                else {
+                    "iso": lambda x: DateTimeFieldV2.__fromisoformat_py310_fix("iso", x),
+                    "iso8601": lambda x: DateTimeFieldV2.__fromisoformat_py310_fix("iso8601", x),
+                }
+            ),
+        }
+
+        def _serialize(self, value: Any, attr: Any, obj: Any, **kwargs: Any) -> Any:
+            if value is None:
+                return None
+
+            if value.tzinfo is None:
+                value = value.replace(tzinfo=datetime.timezone.utc)
+
+            return super()._serialize(value, attr, obj, **kwargs)
+
         def _deserialize(self, value: Any, attr: Any, data: Any, **_: Any) -> Any:
             result = super()._deserialize(value, attr, data)
             if result.tzinfo is None:
                 return result.replace(tzinfo=datetime.timezone.utc)
+            if result.tzinfo == datetime.timezone.utc:
+                return result
             if dateutil_tz_utc_cls is not None and isinstance(result.tzinfo, dateutil_tz_utc_cls):
                 return result.replace(tzinfo=datetime.timezone.utc)
             return result.astimezone(datetime.timezone.utc)

diff --git a/tests/test_serialization.py b/tests/test_serialization.py
@@ -374,6 +374,7 @@ class BoolContainer:
         ("2022-02-20T11:33:48", datetime.datetime(2022, 2, 20, 11, 33, 48, 0, datetime.timezone.utc)),
         ("2022-02-20T11:33:48.607289Z", datetime.datetime(2022, 2, 20, 11, 33, 48, 607289, datetime.timezone.utc)),
         ("2022-02-20T11:33:48Z", datetime.datetime(2022, 2, 20, 11, 33, 48, 0, datetime.timezone.utc)),
+        ("2022-02-20", datetime.datetime(2022, 2, 20, 0, 0, 0, 0, datetime.timezone.utc)),
     ],
 )
 def test_datetime_field_load(raw: str, dt: datetime.datetime) -> None:
@@ -392,6 +393,7 @@ class DateTimeContainer:
         (datetime.datetime(2022, 2, 20, 11, 33, 48, 0, datetime.timezone.utc), "2022-02-20T11:33:48+00:00"),
         (datetime.datetime(2022, 2, 20, 11, 33, 48, 607289, None), "2022-02-20T11:33:48.607289+00:00"),
         (datetime.datetime(2022, 2, 20, 11, 33, 48, 0, None), "2022-02-20T11:33:48+00:00"),
+        (datetime.datetime(2022, 2, 20, 0, 0, 0, 0, None), "2022-02-20T00:00:00+00:00"),
     ],
 )
 def test_datetime_field_dump(dt: datetime.datetime, raw: str) -> None: