fix the astimezone the replace methods (#189)

Author: frostming

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+- Fix the `astimezone()` and `replace()` methods of datetime objects. ([#188](https://github.com/sdispater/tomlkit/issues/178))
+
 ## [0.10.1] - 2022-03-27
 
 ### Fixed

tests/test_items.py

@@ -8,10 +8,8 @@
 
 import pytest
 
-from tomlkit import inline_table
+from tomlkit import api
 from tomlkit import parse
-from tomlkit.api import array
-from tomlkit.api import ws
 from tomlkit.exceptions import NonExistentKey
 from tomlkit.items import Bool
 from tomlkit.items import Comment
@@ -27,6 +25,50 @@
 from tomlkit.parser import Parser
 
 
+@pytest.fixture()
+def tz_pst():
+    try:
+        from datetime import timezone
+
+        return timezone(timedelta(hours=-8), "PST")
+    except ImportError:
+        from datetime import tzinfo
+
+        class PST(tzinfo):
+            def utcoffset(self, dt):
+                return timedelta(hours=-8)
+
+            def tzname(self, dt):
+                return "PST"
+
+            def dst(self, dt):
+                return timedelta(0)
+
+        return PST()
+
+
+@pytest.fixture()
+def tz_utc():
+    try:
+        from datetime import timezone
+
+        return timezone.utc
+    except ImportError:
+        from datetime import tzinfo
+
+        class UTC(tzinfo):
+            def utcoffset(self, dt):
+                return timedelta(hours=0)
+
+            def tzname(self, dt):
+                return "UTC"
+
+            def dst(self, dt):
+                return timedelta(0)
+
+        return UTC()
+
+
 def test_key_comparison():
     k = Key("foo")
 
@@ -320,10 +362,10 @@ def test_dicts_are_converted_to_tables():
 
 
 def test_array_add_line():
-    t = array()
+    t = api.array()
     t.add_line(1, 2, 3, comment="Line 1")
     t.add_line(4, 5, 6, comment="Line 2")
-    t.add_line(7, ws(","), ws(" "), 8, add_comma=False)
+    t.add_line(7, api.ws(","), api.ws(" "), 8, add_comma=False)
     t.add_line(indent="")
     assert len(t) == 8
     assert list(t) == [1, 2, 3, 4, 5, 6, 7, 8]
@@ -338,9 +380,9 @@ def test_array_add_line():
 
 
 def test_array_add_line_invalid_value():
-    t = array()
+    t = api.array()
     with pytest.raises(ValueError, match="is not allowed"):
-        t.add_line(1, ws(" "))
+        t.add_line(1, api.ws(" "))
     with pytest.raises(ValueError, match="is not allowed"):
         t.add_line(Comment(Trivia("  ", comment="test")))
     assert len(t) == 0
@@ -465,7 +507,7 @@ def test_floats_behave_like_floats():
     assert doc.as_string() == "float = +35.12"
 
 
-def test_datetimes_behave_like_datetimes():
+def test_datetimes_behave_like_datetimes(tz_utc, tz_pst):
     i = item(datetime(2018, 7, 22, 12, 34, 56))
 
     assert i == datetime(2018, 7, 22, 12, 34, 56)
@@ -479,6 +521,14 @@ def test_datetimes_behave_like_datetimes():
     assert i == datetime(2018, 7, 21, 12, 34, 56)
     assert i.as_string() == "2018-07-21T12:34:56"
 
+    i = i.replace(year=2019, tzinfo=tz_utc)
+    assert i == datetime(2019, 7, 21, 12, 34, 56, tzinfo=tz_utc)
+    assert i.as_string() == "2019-07-21T12:34:56+00:00"
+
+    i = i.astimezone(tz_pst)
+    assert i == datetime(2019, 7, 21, 4, 34, 56, tzinfo=tz_pst)
+    assert i.as_string() == "2019-07-21T04:34:56-08:00"
+
     doc = parse("dt = 2018-07-22T12:34:56-05:00")
     doc["dt"] += timedelta(days=1)
 
@@ -499,6 +549,10 @@ def test_dates_behave_like_dates():
     assert i == date(2018, 7, 21)
     assert i.as_string() == "2018-07-21"
 
+    i = i.replace(year=2019)
+    assert i == datetime(2019, 7, 21)
+    assert i.as_string() == "2019-07-21"
+
     doc = parse("dt = 2018-07-22 # Comment")
     doc["dt"] += timedelta(days=1)
 
@@ -511,6 +565,10 @@ def test_times_behave_like_times():
     assert i == time(12, 34, 56)
     assert i.as_string() == "12:34:56"
 
+    i = i.replace(hour=13)
+    assert i == time(13, 34, 56)
+    assert i.as_string() == "13:34:56"
+
 
 def test_strings_behave_like_strs():
     i = item("foo")
@@ -611,7 +669,7 @@ def test_items_are_pickable():
     s = pickle.dumps(n)
     assert pickle.loads(s).as_string() == 'foo = "bar"\n'
 
-    n = inline_table()
+    n = api.inline_table()
     n["foo"] = "bar"
 
     s = pickle.dumps(n)
@@ -629,7 +687,7 @@ def test_items_are_pickable():
 
 
 def test_trim_comments_when_building_inline_table():
-    table = inline_table()
+    table = api.inline_table()
     row = parse('foo = "bar"  # Comment')
     table.update(row)
     assert table.as_string() == '{foo = "bar"}'
@@ -641,7 +699,7 @@ def test_trim_comments_when_building_inline_table():
 
 
 def test_deleting_inline_table_elemeent_does_not_leave_trailing_separator():
-    table = inline_table()
+    table = api.inline_table()
     table["foo"] = "bar"
     table["baz"] = "boom"
 
@@ -651,7 +709,7 @@ def test_deleting_inline_table_elemeent_does_not_leave_trailing_separator():
 
     assert '{foo = "bar"}' == table.as_string()
 
-    table = inline_table()
+    table = api.inline_table()
     table["foo"] = "bar"
 
     del table["foo"]

tomlkit/items.py

@@ -721,8 +721,7 @@ def __new__(
         second: int,
         microsecond: int,
         tzinfo: Optional[tzinfo],
-        trivia: Trivia,
-        raw: str,
+        *_: Any,
         **kwargs: Any,
     ) -> datetime:
         return datetime.__new__(
@@ -748,12 +747,13 @@ def __init__(
         second: int,
         microsecond: int,
         tzinfo: Optional[tzinfo],
-        trivia: Trivia,
-        raw: str,
+        trivia: Optional[Trivia] = None,
+        raw: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
-        super().__init__(trivia)
+        super().__init__(trivia or Trivia())
 
-        self._raw = raw
+        self._raw = raw or self.isoformat()
 
     @property
     def discriminant(self) -> int:
@@ -803,7 +803,16 @@ def __sub__(self, other):
 
         return result
 
-    def _new(self, result):
+    def replace(self, *args: Any, **kwargs: Any) -> datetime:
+        return self._new(super().replace(*args, **kwargs))
+
+    def astimezone(self, tz: tzinfo) -> datetime:
+        result = super().astimezone(tz)
+        if PY38:
+            return result
+        return self._new(result)
+
+    def _new(self, result) -> "DateTime":
         raw = result.isoformat()
 
         return DateTime(
@@ -879,6 +888,9 @@ def __sub__(self, other):
 
         return result
 
+    def replace(self, *args: Any, **kwargs: Any) -> date:
+        return self._new(super().replace(*args, **kwargs))
+
     def _new(self, result):
         raw = result.isoformat()
 
@@ -929,6 +941,22 @@ def value(self) -> time:
     def as_string(self) -> str:
         return self._raw
 
+    def replace(self, *args: Any, **kwargs: Any) -> time:
+        return self._new(super().replace(*args, **kwargs))
+
+    def _new(self, result):
+        raw = result.isoformat()
+
+        return Time(
+            result.hour,
+            result.minute,
+            result.second,
+            result.microsecond,
+            result.tzinfo,
+            self._trivia,
+            raw,
+        )
+
     def _getstate(self, protocol: int = 3) -> tuple:
         return (
             self.hour,