Skip to content

refactor: Simplify calculate_timestamp_datetime and calculate_timestamp_date #2838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FBruzzesi merged 4 commits into from
Jul 17, 2025
Merged

refactor: Simplify calculate_timestamp_datetime and calculate_timestamp_date #2838

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0e6def9
refactor: Simplify calculate_timestamp_datetime and calculate_timesta…
FBruzzesi Jul 16, 2025
bdb46aa
fix(typing): Remove `NativeSeriesT` from alias def
dangotbanned Jul 17, 2025
eb8ec7f
Merge remote-tracking branch 'upstream/main' into refactor/simplify-c…
dangotbanned Jul 17, 2025
86efcc7
Merge branch 'main' into refactor/simplify-calculate_timestamp_datetime
dangotbanned Jul 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 43 additions & 43 deletions narwhals/_pandas_like/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import functools
import operator
import re
from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar

Expand Down Expand Up @@ -44,6 +45,10 @@
from narwhals.typing import DTypeBackend, IntoDType, TimeUnit, _1DArray

ExprT = TypeVar("ExprT", bound=PandasLikeExpr)
UnitCurrent: TypeAlias = TimeUnit
UnitTarget: TypeAlias = TimeUnit
BinOpBroadcast: TypeAlias = Callable[[Any, int], Any]
IntoRhs: TypeAlias = int


PANDAS_LIKE_IMPLEMENTATION = {
Expand Down Expand Up @@ -553,52 +558,47 @@ def int_dtype_mapper(dtype: Any) -> str:
return "int64"


def calculate_timestamp_datetime( # noqa: C901, PLR0912
s: NativeSeriesT, original_time_unit: str, time_unit: str
_TIMESTAMP_DATETIME_OP_FACTOR: Mapping[
tuple[UnitCurrent, UnitTarget], tuple[BinOpBroadcast, IntoRhs]
Copy link
Member

@dangotbanned dangotbanned Jul 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#2838 (comment)

I think the issue is that BinOpBroadcast is generic, but using a TypeVar with a default.

So you don't get a warning when using BinOpBroadcast without specialising - but it causes an issue later when the default conflicts with the constraints of the function signature.

Hopefully this works, I'm replying from my phone 😅

Suggested change
tuple[UnitCurrent, UnitTarget], tuple[BinOpBroadcast, IntoRhs]
tuple[UnitCurrent, UnitTarget], tuple[BinOpBroadcast[NativeSeriesT], IntoRhs]

Copy link
Member

@dangotbanned dangotbanned Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

narwhals/_pandas_like/utils.py:583: error: Argument 1 has incompatible type "_CuDFSeries"; expected "Series[Any]"  [arg-type]
narwhals/_pandas_like/utils.py:583: error: Argument 1 has incompatible type "_ModinSeries"; expected "Series[Any]"  [arg-type]

narwhals/narwhals/_pandas_like/typing.py

Lines 25 to 31 in bb03703

NativeSeriesT = TypeVar(
"NativeSeriesT",
"pd.Series[Any]",
"_CuDFSeries",
"_ModinSeries",
default="pd.Series[Any]",
)

NativeSeriesT means exactly one of pd.Series[Any], _CuDFSeries, _ModinSeries, but default to pd.Series[Any] when not provided.

So the conflict is that

  • NativeSeriesT is still used for the function parameter and return type
  • But _TIMESTAMP_DATETIME_OP_FACTOR is currently saying it only accepts pd.Series[Any]

narwhals/narwhals/_pandas_like/utils.py

Lines 576 to 578 in 0e6def9

def calculate_timestamp_datetime(
s: NativeSeriesT, current: TimeUnit, time_unit: TimeUnit
) -> NativeSeriesT:

Copy link
Member Author

@FBruzzesi FBruzzesi Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @dangotbanned thanks for investigating this.

But _TIMESTAMP_DATETIME_OP_FACTOR is currently saying it only accepts pd.Series[Any]

Why would that be the case? The definition is:

BinOpBroadcast: TypeAlias = Callable[[NativeSeriesT, int], NativeSeriesT]

Regarding #2838 (comment), BinOpBroadcast is not a generic, and it would result in:

Type variable "NativeSeriesT" has no meaning in this contextPylance[reportGeneralTypeIssues]

🧐

Copy link
Member

@dangotbanned dangotbanned Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @dangotbanned thanks for investigating this.

But _TIMESTAMP_DATETIME_OP_FACTOR is currently saying it only accepts pd.Series[Any]

Why would that be the case? The definition is:

BinOpBroadcast: TypeAlias = Callable[[NativeSeriesT, int], NativeSeriesT]

Regarding #2838 (comment), BinOpBroadcast is not a generic, and it would result in:

BinOpBroadcast is a generic TypeAlias because it uses a TypeVar (NativeSeriesT) in it's definition.

image

If we compare that to when you hover over the current use, it uses the default type - because it wasn't subscribed

image

Type variable "NativeSeriesT" has no meaning in this contextPylance[reportGeneralTypeIssues]

🧐

Ahh, so that message it about scoping - not that the TypeVar can't be used at all

So this is because _TIMESTAMP_DATETIME_OP_FACTOR is in a global scope, but we're trying to solve the TypeVar in a function scope

Copy link
Member

@dangotbanned dangotbanned Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've just used Any instead in (bdb46aa)

The main reason I had that typing in _arrow was because the stubs caused too many false positives when I used the functions directly.

Here, we're just using operator functions, which are mostly typed with Any anyway 😉

https://github.com/python/typeshed/blob/8e25bda401332ac56d4750cffda2127a4302cd2f/stdlib/_operator.pyi#L42-L71

Copy link
Member

@dangotbanned dangotbanned Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We still have NativeSeriesT preserved on the outside, so this is simpler than refactoring everything to provide a scope for the TypeVar 😅

] = {
("ns", "us"): (operator.floordiv, 1_000),
("ns", "ms"): (operator.floordiv, 1_000_000),
("us", "ns"): (operator.mul, NS_PER_MICROSECOND),
("us", "ms"): (operator.floordiv, 1_000),
("ms", "ns"): (operator.mul, NS_PER_MILLISECOND),
("ms", "us"): (operator.mul, 1_000),
("s", "ns"): (operator.mul, NS_PER_SECOND),
("s", "us"): (operator.mul, US_PER_SECOND),
("s", "ms"): (operator.mul, MS_PER_SECOND),
}


def calculate_timestamp_datetime(
s: NativeSeriesT, current: TimeUnit, time_unit: TimeUnit
) -> NativeSeriesT:
if original_time_unit == "ns":
if time_unit == "ns":
result = s
elif time_unit == "us":
result = s // 1_000
else:
result = s // 1_000_000
elif original_time_unit == "us":
if time_unit == "ns":
result = s * NS_PER_MICROSECOND
elif time_unit == "us":
result = s
else:
result = s // 1_000
elif original_time_unit == "ms":
if time_unit == "ns":
result = s * NS_PER_MILLISECOND
elif time_unit == "us":
result = s * 1_000
else:
result = s
elif original_time_unit == "s":
if time_unit == "ns":
result = s * NS_PER_SECOND
elif time_unit == "us":
result = s * US_PER_SECOND
else:
result = s * MS_PER_SECOND
if current == time_unit:
return s
elif item := _TIMESTAMP_DATETIME_OP_FACTOR.get((current, time_unit)):
fn, factor = item
return fn(s, factor)
else: # pragma: no cover
msg = f"unexpected time unit {original_time_unit}, please report a bug at https://github.com/narwhals-dev/narwhals"
msg = (
f"unexpected time unit {current}, please report an issue at "
"https://github.com/narwhals-dev/narwhals"
)
raise AssertionError(msg)
return result


def calculate_timestamp_date(s: NativeSeriesT, time_unit: str) -> NativeSeriesT:
s = s * SECONDS_PER_DAY
if time_unit == "ns":
result = s * NS_PER_SECOND
elif time_unit == "us":
result = s * US_PER_SECOND
else:
result = s * MS_PER_SECOND
return result


_TIMESTAMP_DATE_FACTOR: Mapping[TimeUnit, int] = {
"ns": NS_PER_SECOND,
"us": US_PER_SECOND,
"ms": MS_PER_SECOND,
"s": 1,
}


def calculate_timestamp_date(s: NativeSeriesT, time_unit: TimeUnit) -> NativeSeriesT:
return s * SECONDS_PER_DAY * _TIMESTAMP_DATE_FACTOR[time_unit]


def select_columns_by_name(
Expand Down
Loading