Preserve temporal fold during copying and pickling

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Fixed
 
+- Preserve the `fold` of datetime and time items during copying and pickle protocols 4 and later, matching the standard-library types.
+
 - Fix valid TOML being rejected when an out-of-order table reopens an implicit super-table with a dotted key: `[a.b.c]` then `[a]` with `b.z = 1` extends the implicit `a.b` super-table rather than redefining it (stdlib `tomllib` accepts it), so it now parses and round-trips instead of raising `Redefinition of an existing table`. The concrete-table case (`[a.b]` then `[a]` `b.z = 1`) stays rejected. ([#565](https://github.com/python-poetry/tomlkit/pull/565))
 - Fix appending a key to a parsed inline table that uses `", "` separators emitting a bare `","` for the new entry, leaving the table inconsistently spaced (e.g. `{x = 1, y = 2,z = 3}`). ([#595](https://github.com/python-poetry/tomlkit/pull/595))
 

tests/test_items.py

@@ -86,18 +86,19 @@ def dst(self, dt: datetime | None) -> timedelta:
         return UTC()
 
 
-@pytest.fixture()
-def tz_fold() -> tzinfo:
-    class FoldTimezone(tzinfo):
-        def utcoffset(self, dt: datetime | None) -> timedelta:
-            return timedelta(hours=-5 if dt is not None and dt.fold else -4)
+class FoldTimezone(tzinfo):
+    def utcoffset(self, dt: datetime | None) -> timedelta:
+        return timedelta(hours=-5 if dt is not None and dt.fold else -4)
 
-        def tzname(self, dt: datetime | None) -> str:
-            return "FoldTimezone"
+    def tzname(self, dt: datetime | None) -> str:
+        return "FoldTimezone"
 
-        def dst(self, dt: datetime | None) -> timedelta:
-            return timedelta(0)
+    def dst(self, dt: datetime | None) -> timedelta:
+        return timedelta(0)
 
+
+@pytest.fixture()
+def tz_fold() -> tzinfo:
     return FoldTimezone()
 
 
@@ -1094,6 +1095,55 @@ def test_items_are_pickable() -> None:
     assert pickle.loads(s).as_string() == 'foo = "bar"\n'
 
 
+@pytest.mark.parametrize("fold", [0, 1])
+@pytest.mark.parametrize("kind", ["datetime", "time"])
+@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1))
+def test_temporal_pickle_preserves_fold(
+    tz_fold: tzinfo, fold: int, kind: str, protocol: int
+) -> None:
+    value = (
+        datetime(2018, 11, 4, 1, 30, tzinfo=tz_fold, fold=fold)
+        if kind == "datetime"
+        else time(1, 30, fold=fold)
+    )
+    wrapped = item(value).comment("ambiguous time")
+    assert isinstance(wrapped, (DateTime, Time))
+    result = pickle.loads(pickle.dumps(wrapped, protocol=protocol))
+    expected = pickle.loads(pickle.dumps(value, protocol=protocol))
+
+    assert type(result) is type(wrapped)
+    assert result.fold == expected.fold
+    assert result.as_string() == wrapped.as_string()
+    assert result.trivia == wrapped.trivia
+    if isinstance(expected, datetime):
+        assert result.timestamp() == expected.timestamp()
+
+
+@pytest.mark.parametrize("fold", [0, 1])
+@pytest.mark.parametrize("kind", ["datetime", "time"])
+@pytest.mark.parametrize("deep", [False, True])
+def test_temporal_copy_preserves_fold(
+    tz_fold: tzinfo, fold: int, kind: str, deep: bool
+) -> None:
+    value = (
+        datetime(2018, 11, 4, 1, 30, tzinfo=tz_fold, fold=fold)
+        if kind == "datetime"
+        else time(1, 30, fold=fold)
+    )
+    wrapped = item(value).comment("ambiguous time")
+    assert isinstance(wrapped, (DateTime, Time))
+    result = copy.deepcopy(wrapped) if deep else copy.copy(wrapped)
+
+    assert type(result) is type(wrapped)
+    assert result is not wrapped
+    assert result.fold == value.fold
+    assert result.as_string() == wrapped.as_string()
+    assert result.trivia == wrapped.trivia
+    if isinstance(value, datetime):
+        assert isinstance(result, DateTime)
+        assert result.timestamp() == value.timestamp()
+
+
 def test_trim_comments_when_building_inline_table() -> None:
     table = api.inline_table()
     row = parse('foo = "bar"  # Comment')

tomlkit/items.py

@@ -1204,6 +1204,12 @@ def _new(self, result: datetime) -> DateTime:
             fold=result.fold,
         )
 
+    def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[object, ...]]:  # type: ignore[override]
+        cls, args = super().__reduce_ex__(protocol)
+        if protocol >= 4 and self.fold:
+            args += (self.fold,)
+        return cls, args
+
     def _getstate(
         self, protocol: int = 3
     ) -> tuple[int, int, int, int, int, int, int, tzinfo | None, Trivia, str]:
@@ -1365,6 +1371,12 @@ def _new(self, result: time) -> Time:
             fold=result.fold,
         )
 
+    def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[object, ...]]:  # type: ignore[override]
+        cls, args = super().__reduce_ex__(protocol)
+        if protocol >= 4 and self.fold:
+            args += (self.fold,)
+        return cls, args
+
     def _getstate(
         self, protocol: int = 3
     ) -> tuple[int, int, int, int, tzinfo | None, Trivia, str]: