Overhaul (deep)copy and pickle support for temporal and spatial types (#1002)

* Add support to pickling DateTime

* Fix Duration copy and deepcopy dropping dynamic attributes (in `__dict__`)

* Adding more unit tests

---------

Co-authored-by: Robsdedude <[email protected]>

Author: Tanguy-Boisset

src/neo4j/time/__init__.py

@@ -574,12 +574,17 @@ def __str__(self) -> str:
         """"""
         return self.iso_format()
 
-    def __copy__(self) -> Duration:
-        return self.__new__(self.__class__, months=self[0], days=self[1],
-                            seconds=self[2], nanoseconds=self[3])
+    def __reduce__(self):
+        return (
+            type(self)._restore, (tuple(self), self.__dict__)
+        )
 
-    def __deepcopy__(self, memo) -> Duration:
-        return self.__copy__()
+    @classmethod
+    def _restore(cls, elements, dict_):
+        instance = tuple.__new__(cls, elements)
+        if dict_:
+            instance.__dict__.update(dict_)
+        return instance
 
     @classmethod
     def from_iso_format(cls, s: str) -> Duration:
@@ -763,6 +768,13 @@ class Date(date_base_class, metaclass=DateType):
     # CONSTRUCTOR #
 
     def __new__(cls, year: int, month: int, day: int) -> Date:
+        # TODO: 6.0 - remove the __new__ magic and ZeroDate being a singleton.
+        #             It's fine to remain as constant. Instead, simply use
+        #             __init__ and simplify pickle/copy (remove __reduce__).
+        #             N.B. this is a breaking change and must be treated as
+        #             such. Also consider introducing __slots__. Potentially
+        #             apply similar treatment to other temporal types as well
+        #             as spatial types.
         if year == month == day == 0:
             return ZeroDate
         year, month, day = _normalize_day(year, month, day)
@@ -1218,11 +1230,17 @@ def __sub__(self, other):
         except TypeError:
             return NotImplemented
 
-    def __copy__(self) -> Date:
-        return self.__new(self.__ordinal, self.__year, self.__month, self.__day)
+    def __reduce__(self):
+        if self is ZeroDate:
+            return "ZeroDate"
+        return type(self)._restore, (self.__dict__,)
 
-    def __deepcopy__(self, *args, **kwargs) -> Date:
-        return self.__copy__()
+    @classmethod
+    def _restore(cls, dict_) -> Date:
+        instance = object.__new__(cls)
+        if dict_:
+            instance.__dict__.update(dict_)
+        return instance
 
     # INSTANCE METHODS #
 
@@ -1396,34 +1414,53 @@ class Time(time_base_class, metaclass=TimeType):
 
     # CONSTRUCTOR #
 
-    def __new__(
-        cls,
+    def __init__(
+        self,
         hour: int = 0,
         minute: int = 0,
         second: int = 0,
         nanosecond: int = 0,
         tzinfo: t.Optional[_tzinfo] = None
-    ) -> Time:
-        hour, minute, second, nanosecond = cls.__normalize_nanosecond(
+    ) -> None:
+        hour, minute, second, nanosecond = self.__normalize_nanosecond(
             hour, minute, second, nanosecond
         )
         ticks = (3600000000000 * hour
                  + 60000000000 * minute
                  + 1000000000 * second
                  + nanosecond)
-        return cls.__new(ticks, hour, minute, second, nanosecond, tzinfo)
+        self.__unchecked_init(ticks, hour, minute, second, nanosecond, tzinfo)
 
     @classmethod
-    def __new(cls, ticks, hour, minute, second, nanosecond, tzinfo):
-        instance = object.__new__(cls)
-        instance.__ticks = int(ticks)
-        instance.__hour = int(hour)
-        instance.__minute = int(minute)
-        instance.__second = int(second)
-        instance.__nanosecond = int(nanosecond)
-        instance.__tzinfo = tzinfo
+    def __unchecked_new(
+        cls,
+        ticks: int,
+        hour: int,
+        minutes: int,
+        second: int,
+        nano: int,
+        tz: t.Optional[_tzinfo]
+    ) -> Time:
+        instance = object.__new__(Time)
+        instance.__unchecked_init(ticks, hour, minutes, second, nano, tz)
         return instance
 
+    def __unchecked_init(
+        self,
+        ticks: int,
+        hour: int,
+        minutes: int,
+        second: int,
+        nano: int,
+        tz: t.Optional[_tzinfo]
+    ) -> None:
+        self.__ticks = ticks
+        self.__hour = hour
+        self.__minute = minutes
+        self.__second = second
+        self.__nanosecond = nano
+        self.__tzinfo = tz
+
     # CLASS METHODS #
 
     @classmethod
@@ -1521,7 +1558,8 @@ def from_ticks(cls, ticks: int, tz: t.Optional[_tzinfo] = None) -> Time:
             second, nanosecond = divmod(ticks, NANO_SECONDS)
             minute, second = divmod(second, 60)
             hour, minute = divmod(minute, 60)
-            return cls.__new(ticks, hour, minute, second, nanosecond, tz)
+            return cls.__unchecked_new(ticks, hour, minute, second, nanosecond,
+                                       tz)
         raise ValueError("Ticks out of range (0..86400000000000)")
 
     @classmethod
@@ -1619,7 +1657,7 @@ def utc_now(cls) -> Time:
 
     __nanosecond = 0
 
-    __tzinfo = None
+    __tzinfo: t.Optional[_tzinfo] = None
 
     @property
     def ticks(self) -> int:
@@ -1751,13 +1789,6 @@ def __gt__(self, other: t.Union[Time, time]) -> bool:
             return NotImplemented
         return self_ticks > other_ticks
 
-    def __copy__(self) -> Time:
-        return self.__new(self.__ticks, self.__hour, self.__minute,
-                          self.__second, self.__nanosecond, self.__tzinfo)
-
-    def __deepcopy__(self, *args, **kwargs) -> Time:
-        return self.__copy__()
-
     # INSTANCE METHODS #
 
     if t.TYPE_CHECKING:
@@ -2126,6 +2157,10 @@ def combine(  # type: ignore[override]
         """
         assert isinstance(date, Date)
         assert isinstance(time, Time)
+        return cls._combine(date, time)
+
+    @classmethod
+    def _combine(cls, date: Date, time: Time) -> DateTime:
         instance = object.__new__(cls)
         instance.__date = date
         instance.__time = time
@@ -2491,11 +2526,15 @@ def __sub__(self, other):
             return self.__add__(-other)
         return NotImplemented
 
-    def __copy__(self) -> DateTime:
-        return self.combine(self.__date, self.__time)
+    def __reduce__(self):
+        return type(self)._restore, (self.__dict__,)
 
-    def __deepcopy__(self, memo) -> DateTime:
-        return self.__copy__()
+    @classmethod
+    def _restore(cls, dict_):
+        instance = object.__new__(cls)
+        if dict_:
+            instance.__dict__.update(dict_)
+        return instance
 
     # INSTANCE METHODS #
 

tests/unit/common/spatial/test_cartesian_point.py

@@ -16,11 +16,23 @@
 
 from __future__ import annotations
 
+import copy
+import pickle
+
 import pytest
 
 from neo4j.spatial import CartesianPoint
 
 
+def make_reduce_points():
+    return (
+        CartesianPoint((1, 2)),
+        CartesianPoint((3.2, 4.0)),
+        CartesianPoint((3, 4, -1)),
+        CartesianPoint((3.2, 4.0, -1.2)),
+    )
+
+
 class TestCartesianPoint:
 
     def test_alias_3d(self) -> None:
@@ -42,3 +54,30 @@ def test_alias_2d(self) -> None:
         assert p.y == y
         with pytest.raises(AttributeError):
             _ = p.z
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_copy(self, p):
+        p.foo = [1, 2]
+        p2 = copy.copy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo is p2.foo
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_deep_copy(self, p):
+        p.foo = [1, [2]]
+        p2 = copy.deepcopy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo == p2.foo
+        assert p.foo is not p2.foo
+        assert p.foo[1] is not p2.foo[1]
+
+    @pytest.mark.parametrize("expected", make_reduce_points())
+    def test_pickle(self, expected):
+        expected.foo = [1, [2]]
+        actual = pickle.loads(pickle.dumps(expected))
+        assert expected == actual
+        assert expected is not actual
+        assert expected.foo == actual.foo
+        assert expected.foo is not actual.foo

tests/unit/common/spatial/test_point.py

@@ -16,6 +16,8 @@
 
 from __future__ import annotations
 
+import copy
+import pickle
 import typing as t
 
 import pytest
@@ -26,6 +28,17 @@
 )
 
 
+def make_reduce_points():
+    return (
+        Point((42,)),
+        Point((69.420,)),
+        Point((1, 2)),
+        Point((1.2, 2.3)),
+        Point((1, 3, 3, 7)),
+        Point((1.0, 3.0, 3.0, 7.0)),
+    )
+
+
 class TestPoint:
 
     @pytest.mark.parametrize("argument", (
@@ -59,3 +72,30 @@ def test_immutable_coordinates(self) -> None:
             p[1] = 2.0  # type: ignore[index]
         with pytest.raises(TypeError):
             p[2] = 2.0  # type: ignore[index]
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_copy(self, p):
+        p.foo = [1, 2]
+        p2 = copy.copy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo is p2.foo
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_deep_copy(self, p):
+        p.foo = [1, [2]]
+        p2 = copy.deepcopy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo == p2.foo
+        assert p.foo is not p2.foo
+        assert p.foo[1] is not p2.foo[1]
+
+    @pytest.mark.parametrize("expected", make_reduce_points())
+    def test_pickle(self, expected):
+        expected.foo = [1, [2]]
+        actual = pickle.loads(pickle.dumps(expected))
+        assert expected == actual
+        assert expected is not actual
+        assert expected.foo == actual.foo
+        assert expected.foo is not actual.foo

tests/unit/common/spatial/test_wgs84_point.py

@@ -16,11 +16,23 @@
 
 from __future__ import annotations
 
+import copy
+import pickle
+
 import pytest
 
 from neo4j.spatial import WGS84Point
 
 
+def make_reduce_points():
+    return (
+        WGS84Point((1, 2)),
+        WGS84Point((3.2, 4.0)),
+        WGS84Point((3, 4, -1)),
+        WGS84Point((3.2, 4.0, -1.2)),
+    )
+
+
 class TestWGS84Point:
 
     def test_alias_3d(self) -> None:
@@ -60,3 +72,30 @@ def test_alias_2d(self) -> None:
             p.height
         with pytest.raises(AttributeError):
             p.z
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_copy(self, p):
+        p.foo = [1, 2]
+        p2 = copy.copy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo is p2.foo
+
+    @pytest.mark.parametrize("p", make_reduce_points())
+    def test_deep_copy(self, p):
+        p.foo = [1, [2]]
+        p2 = copy.deepcopy(p)
+        assert p == p2
+        assert p is not p2
+        assert p.foo == p2.foo
+        assert p.foo is not p2.foo
+        assert p.foo[1] is not p2.foo[1]
+
+    @pytest.mark.parametrize("expected", make_reduce_points())
+    def test_pickle(self, expected):
+        expected.foo = [1, [2]]
+        actual = pickle.loads(pickle.dumps(expected))
+        assert expected == actual
+        assert expected is not actual
+        assert expected.foo == actual.foo
+        assert expected.foo is not actual.foo