From c0f44e20240771780055408822f6ae7308741d34 Mon Sep 17 00:00:00 2001 From: xiehuanyi <74809554+xiehuanyi@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:26:12 +0300 Subject: [PATCH 1/2] Preserve temporal fold during copying and pickling --- CHANGELOG.md | 2 ++ tests/test_items.py | 68 +++++++++++++++++++++++++++++++++++++++------ tomlkit/items.py | 12 ++++++++ 3 files changed, 73 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index edea7022..80a6c9af 100644 --- a/CHANGELOG.md +++ b/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 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)) ## [0.15.1] - 2026-07-17 diff --git a/tests/test_items.py b/tests/test_items.py index 8a946f9f..e031d4cb 100644 --- a/tests/test_items.py +++ b/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() @@ -1050,6 +1051,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') diff --git a/tomlkit/items.py b/tomlkit/items.py index 07604948..f2f8b8a3 100644 --- a/tomlkit/items.py +++ b/tomlkit/items.py @@ -1185,6 +1185,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: + args += (self.fold,) + return cls, args + def _getstate( self, protocol: int = 3 ) -> tuple[int, int, int, int, int, int, int, tzinfo | None, Trivia, str]: @@ -1346,6 +1352,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: + args += (self.fold,) + return cls, args + def _getstate( self, protocol: int = 3 ) -> tuple[int, int, int, int, tzinfo | None, Trivia, str]: From e0dc99006683479fb56c2730675c1999192f9176 Mon Sep 17 00:00:00 2001 From: xiehuanyi <74809554+xiehuanyi@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:30:16 +0300 Subject: [PATCH 2/2] Keep the original pickle representation for default fold --- tomlkit/items.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tomlkit/items.py b/tomlkit/items.py index f2f8b8a3..90948ced 100644 --- a/tomlkit/items.py +++ b/tomlkit/items.py @@ -1187,7 +1187,7 @@ def _new(self, result: datetime) -> DateTime: def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[object, ...]]: # type: ignore[override] cls, args = super().__reduce_ex__(protocol) - if protocol >= 4: + if protocol >= 4 and self.fold: args += (self.fold,) return cls, args @@ -1354,7 +1354,7 @@ def _new(self, result: time) -> Time: def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[object, ...]]: # type: ignore[override] cls, args = super().__reduce_ex__(protocol) - if protocol >= 4: + if protocol >= 4 and self.fold: args += (self.fold,) return cls, args