Skip to content

Commit 69336bc

Browse files
committed
Fix bugs in neo4j.time.DateTime handling
* Fix `DateTime` +/- `Duration` computation being wildly off by considering the days of the `DateTime` since UNIX epoch twice. * Fix `Result.to_df` not correctly converting `DateTime`s with `tzinfo` if `parse_dates=True` is passed in. * Fix `DateTime.__ne__` (inequality operator) always falling back to `object.__ne__` (comparison by `id`).
1 parent 979480d commit 69336bc

File tree

6 files changed

+355
-39
lines changed

6 files changed

+355
-39
lines changed

src/neo4j/_async/work/result.py

+13-7
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,7 @@ async def to_df(
892892
consumed.
893893
"""
894894
import pandas as pd # type: ignore[import]
895+
import pytz
895896

896897
if not expand:
897898
df = pd.DataFrame(await self.values(), columns=self._keys)
@@ -931,14 +932,19 @@ async def to_df(
931932
).all()
932933
)
933934
]
935+
936+
def datetime_to_timestamp(x):
937+
if not x:
938+
return pd.NaT
939+
tzinfo = getattr(x, "tzinfo", None)
940+
if tzinfo is None:
941+
return pd.Timestamp(x.iso_format())
942+
return pd.Timestamp(
943+
x.astimezone(pytz.UTC).iso_format()
944+
).astimezone(tzinfo)
945+
934946
df[dt_columns] = df[dt_columns].apply(
935-
lambda col: col.map(
936-
lambda x: pd.Timestamp(x.iso_format()).replace(
937-
tzinfo=getattr(x, "tzinfo", None)
938-
)
939-
if x
940-
else pd.NaT
941-
)
947+
lambda col: col.map(datetime_to_timestamp)
942948
)
943949
return df
944950

src/neo4j/_sync/work/result.py

+13-7
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,7 @@ def to_df(
892892
consumed.
893893
"""
894894
import pandas as pd # type: ignore[import]
895+
import pytz
895896

896897
if not expand:
897898
df = pd.DataFrame(self.values(), columns=self._keys)
@@ -931,14 +932,19 @@ def to_df(
931932
).all()
932933
)
933934
]
935+
936+
def datetime_to_timestamp(x):
937+
if not x:
938+
return pd.NaT
939+
tzinfo = getattr(x, "tzinfo", None)
940+
if tzinfo is None:
941+
return pd.Timestamp(x.iso_format())
942+
return pd.Timestamp(
943+
x.astimezone(pytz.UTC).iso_format()
944+
).astimezone(tzinfo)
945+
934946
df[dt_columns] = df[dt_columns].apply(
935-
lambda col: col.map(
936-
lambda x: pd.Timestamp(x.iso_format()).replace(
937-
tzinfo=getattr(x, "tzinfo", None)
938-
)
939-
if x
940-
else pd.NaT
941-
)
947+
lambda col: col.map(datetime_to_timestamp)
942948
)
943949
return df
944950

src/neo4j/time/__init__.py

+8-3
Original file line numberDiff line numberDiff line change
@@ -2560,7 +2560,7 @@ def __ne__(self, other: object) -> bool:
25602560
25612561
Accepts :class:`.DateTime` and :class:`datetime.datetime`.
25622562
"""
2563-
if not isinstance(other, (Date, date)):
2563+
if not isinstance(other, (DateTime, datetime)):
25642564
return NotImplemented
25652565
return not self.__eq__(other)
25662566

@@ -2641,7 +2641,9 @@ def __gt__( # type: ignore[override]
26412641
def __add__(self, other: timedelta | Duration) -> DateTime:
26422642
"""Add a :class:`datetime.timedelta`."""
26432643
if isinstance(other, Duration):
2644-
t = self.to_clock_time() + ClockTime(
2644+
if other == (0, 0, 0, 0):
2645+
return self
2646+
t = self.time().to_clock_time() + ClockTime(
26452647
other.seconds, other.nanoseconds
26462648
)
26472649
days, seconds = symmetric_divmod(t.seconds, 86400)
@@ -2651,8 +2653,11 @@ def __add__(self, other: timedelta | Duration) -> DateTime:
26512653
time_ = Time.from_ticks(seconds * NANO_SECONDS + t.nanoseconds)
26522654
return self.combine(date_, time_).replace(tzinfo=self.tzinfo)
26532655
if isinstance(other, timedelta):
2656+
if other.total_seconds() == 0:
2657+
return self
26542658
t = self.to_clock_time() + ClockTime(
2655-
86400 * other.days + other.seconds, other.microseconds * 1000
2659+
86400 * other.days + other.seconds,
2660+
other.microseconds * 1000,
26562661
)
26572662
days, seconds = symmetric_divmod(t.seconds, 86400)
26582663
date_ = Date.from_ordinal(days + 1)

tests/unit/async_/work/test_result.py

+42-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
from __future__ import annotations
1818

19+
import datetime
1920
import logging
2021
import typing as t
2122
import uuid
@@ -1082,6 +1083,18 @@ async def test_to_df_expand(
10821083
assert df.equals(expected_df)
10831084

10841085

1086+
DTS_AROUND_SWEDISH_DST_CHANGE: tuple[datetime.datetime, ...] = (
1087+
datetime.datetime(2024, 3, 31, 0, 30, 0),
1088+
datetime.datetime(2024, 3, 31, 1, 30, 0),
1089+
datetime.datetime(2024, 3, 31, 2, 30, 0),
1090+
datetime.datetime(2024, 3, 31, 3, 30, 0),
1091+
datetime.datetime(2024, 10, 27, 0, 30, 0),
1092+
datetime.datetime(2024, 10, 27, 1, 30, 0),
1093+
datetime.datetime(2024, 10, 27, 2, 30, 0),
1094+
datetime.datetime(2024, 10, 27, 3, 30, 0),
1095+
)
1096+
1097+
10851098
@pytest.mark.parametrize(
10861099
("keys", "values", "expected_df"),
10871100
(
@@ -1209,7 +1222,7 @@ async def test_to_df_expand(
12091222
columns=["mixed"],
12101223
),
12111224
),
1212-
# Column with only None (should not be transfomred to NaT)
1225+
# Column with only None (should not be transformed to NaT)
12131226
(
12141227
["all_none"],
12151228
[
@@ -1252,6 +1265,34 @@ async def test_to_df_expand(
12521265
columns=["all_none", "mixed", "n"],
12531266
),
12541267
),
1268+
# Difficult timezones
1269+
*(
1270+
(
1271+
["dt_tz"],
1272+
[
1273+
[
1274+
pytz.UTC.localize(
1275+
neo4j_time.DateTime.from_native(dt)
1276+
).astimezone(pytz.timezone("Europe/Stockholm"))
1277+
+ neo4j_time.Duration(nanoseconds=add_ns)
1278+
]
1279+
for dt in DTS_AROUND_SWEDISH_DST_CHANGE
1280+
],
1281+
pd.DataFrame(
1282+
[
1283+
[
1284+
pytz.UTC.localize(pd.Timestamp(dt)).astimezone(
1285+
pytz.timezone("Europe/Stockholm")
1286+
)
1287+
+ pd.Timedelta(add_ns, unit="ns")
1288+
]
1289+
for dt in DTS_AROUND_SWEDISH_DST_CHANGE
1290+
],
1291+
columns=["dt_tz"],
1292+
),
1293+
)
1294+
for add_ns in (0, 1)
1295+
),
12551296
),
12561297
)
12571298
@pytest.mark.parametrize("expand", [True, False])

0 commit comments

Comments
 (0)