diff --git a/src/pendulum/duration.py b/src/pendulum/duration.py index d6cc0657d..0e6c7b76c 100644 --- a/src/pendulum/duration.py +++ b/src/pendulum/duration.py @@ -2,6 +2,7 @@ from datetime import timedelta from typing import TYPE_CHECKING +from typing import SupportsIndex from typing import cast from typing import overload @@ -458,6 +459,29 @@ def __divmod__(self, other: timedelta) -> tuple[int, Duration]: return NotImplemented + def _restore_args(self) -> tuple[int, int, int, int, int, int, int, int, int]: + # Match __deepcopy__: calendar fields live beside the timedelta + # normalization, and timedelta's own reducer drops them. + return ( + self.remaining_days, + self.remaining_seconds, + self.microseconds, + 0, + self.minutes, + self.hours, + self.weeks, + self.years, + self.months, + ) + + def __reduce__(self) -> tuple[type[Self], tuple[int, ...]]: + return self.__reduce_ex__(2) + + def __reduce_ex__( + self, protocol: SupportsIndex + ) -> tuple[type[Self], tuple[int, ...]]: + return self.__class__, self._restore_args() + def __deepcopy__(self, _: dict[int, Self]) -> Self: return self.__class__( days=self.remaining_days, diff --git a/tests/duration/test_behavior.py b/tests/duration/test_behavior.py index cf60282fa..eadb5b70a 100644 --- a/tests/duration/test_behavior.py +++ b/tests/duration/test_behavior.py @@ -2,6 +2,7 @@ import pickle +from copy import copy from copy import deepcopy from datetime import timedelta @@ -20,6 +21,19 @@ def test_pickle() -> None: assert it == it2 +def test_pickle_and_copy_keep_calendar_components() -> None: + original = pendulum.duration(years=2, months=3, weeks=1, days=2) + restored = pickle.loads(pickle.dumps(original)) + copied = copy(original) + + assert (restored.years, restored.months, restored.weeks) == (2, 3, 1) + assert (copied.years, copied.months, copied.weeks) == (2, 3, 1) + assert str(restored) == str(original) + assert str(copied) == str(original) + assert restored == original + assert copied == original + + def test_comparison_to_timedelta() -> None: duration = pendulum.duration(days=3)