Skip to content

Commit

Permalink
fix: np.datetime64 scalar with ns resolution was become int instead o…
Browse files Browse the repository at this point in the history
…f datetime in to_py_scalar
  • Loading branch information
MarcoGorelli committed Nov 2, 2024
1 parent 5c3db5b commit d33e7ef
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 19 deletions.
7 changes: 7 additions & 0 deletions narwhals/translate.py
Original file line number Diff line number Diff line change
Expand Up @@ -857,6 +857,13 @@ def to_py_scalar(scalar_like: Any) -> Any:
return scalar_like

np = get_numpy()
if (
np
and isinstance(scalar_like, np.datetime64)
and scalar_like.dtype == "datetime64[ns]"
):
return datetime(1970, 1, 1) + timedelta(microseconds=scalar_like.item() // 1000)

if np and np.isscalar(scalar_like) and hasattr(scalar_like, "item"):
return scalar_like.item()

Expand Down
27 changes: 8 additions & 19 deletions tests/translate/to_py_scalar_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

from datetime import datetime
from datetime import timedelta
from typing import TYPE_CHECKING
from typing import Any

import numpy as np
Expand All @@ -12,9 +11,6 @@
import narwhals.stable.v1 as nw
from narwhals.dependencies import get_cudf

if TYPE_CHECKING:
from tests.utils import ConstructorEager


@pytest.mark.parametrize(
("input_value", "expected"),
Expand All @@ -28,28 +24,21 @@
(b"a", b"a"),
(datetime(2021, 1, 1), datetime(2021, 1, 1)),
(timedelta(days=1), timedelta(days=1)),
(pd.Timestamp("2020-01-01"), datetime(2020, 1, 1)),
(pd.Timedelta(days=3), timedelta(days=3)),
(np.datetime64("2020-01-01", "s"), datetime(2020, 1, 1)),
(np.datetime64("2020-01-01", "ms"), datetime(2020, 1, 1)),
(np.datetime64("2020-01-01", "us"), datetime(2020, 1, 1)),
(np.datetime64("2020-01-01", "ns"), datetime(2020, 1, 1)),
],
)
def test_to_py_scalar(
constructor_eager: ConstructorEager,
input_value: Any,
expected: Any,
request: pytest.FixtureRequest,
) -> None:
if isinstance(input_value, bytes) and "cudf" in str(constructor_eager):
request.applymarker(pytest.mark.xfail)
df = nw.from_native(constructor_eager({"a": [input_value]}))
output = nw.to_py_scalar(df["a"].item(0))
if expected == 1 and constructor_eager.__name__.startswith("pandas"):
output = nw.to_py_scalar(input_value)
if expected == 1:
assert not isinstance(output, np.int64)
elif isinstance(expected, datetime) and constructor_eager.__name__.startswith(
"pandas"
):
assert not isinstance(output, pd.Timestamp)
elif isinstance(expected, timedelta) and constructor_eager.__name__.startswith(
"pandas"
):
assert not isinstance(output, pd.Timedelta)
assert output == expected


Expand Down

0 comments on commit d33e7ef

Please sign in to comment.