diff --git a/src/pendulum/datetime.py b/src/pendulum/datetime.py index 96aff25b..fa917732 100644 --- a/src/pendulum/datetime.py +++ b/src/pendulum/datetime.py @@ -50,6 +50,22 @@ _MONTH_CALENDAR = calendar.Calendar(firstweekday=calendar.MONDAY) +def _restore_datetime( + cls: type[DateTime], + year: int, + month: int, + day: int, + hour: int, + minute: int, + second: int, + microsecond: int, + tzinfo: datetime.tzinfo | None, + fold: int, +) -> DateTime: + """Rebuild a DateTime from the state produced by ``DateTime._getstate()``.""" + return cls(year, month, day, hour, minute, second, microsecond, tzinfo, fold=fold) + + class DateTime(datetime.datetime, Date): EPOCH: ClassVar[DateTime] min: ClassVar[DateTime] @@ -1367,21 +1383,13 @@ def _getstate( self.tzinfo, ) - def __reduce__( - self, - ) -> tuple[ - type[Self], - tuple[int, int, int, int, int, int, int, datetime.tzinfo | None], - ]: + def __reduce__(self) -> tuple[Callable[..., DateTime], tuple[Any, ...]]: return self.__reduce_ex__(2) def __reduce_ex__( self, protocol: SupportsIndex - ) -> tuple[ - type[Self], - tuple[int, int, int, int, int, int, int, datetime.tzinfo | None], - ]: - return self.__class__, self._getstate(protocol) + ) -> tuple[Callable[..., DateTime], tuple[Any, ...]]: + return _restore_datetime, (self.__class__, *self._getstate(protocol), self.fold) def __deepcopy__(self, _: dict[int, Self]) -> Self: return self.__class__( diff --git a/tests/datetime/test_behavior.py b/tests/datetime/test_behavior.py index 3ccfc421..4e78080a 100644 --- a/tests/datetime/test_behavior.py +++ b/tests/datetime/test_behavior.py @@ -3,6 +3,7 @@ import pickle import zoneinfo +from copy import copy from copy import deepcopy from datetime import date from datetime import datetime @@ -166,6 +167,37 @@ def test_pickle(): assert dt1 == dt2 +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +@pytest.mark.parametrize("fold", [0, 1]) +def test_pickle_preserves_fold_and_instant(protocol, fold): + original = pendulum.datetime(2024, 11, 3, 1, 30, tz="America/Chicago", fold=fold) + restored = pickle.loads(pickle.dumps(original, protocol=protocol)) + + assert type(restored) is type(original) + assert restored.fold == fold + assert restored.utcoffset() == original.utcoffset() + assert restored.timestamp() == original.timestamp() + + +def test_shallow_copy_preserves_fold_and_instant(): + original = pendulum.datetime(2024, 11, 3, 1, 30, tz="America/Chicago", fold=1) + restored = copy(original) + + assert restored.fold == original.fold + assert restored.utcoffset() == original.utcoffset() + assert restored.timestamp() == original.timestamp() + + +def test_pickle_keeps_tzinfo(): + tz = zoneinfo.ZoneInfo("America/Chicago") + original = pendulum.DateTime(2024, 11, 3, 1, 30, tzinfo=tz, fold=1) + restored = pickle.loads(pickle.dumps(original)) + + assert type(restored.tzinfo) is zoneinfo.ZoneInfo + assert restored.utcoffset() == original.utcoffset() + assert restored.timestamp() == original.timestamp() + + def test_pickle_with_integer_tzinfo(): dt1 = pendulum.datetime(2016, 8, 27, 12, 34, 56, 123456, tz=0) s = pickle.dumps(dt1)