Skip to content

fix(typing): Overhaul @overloads in nw.from_native #3125

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

Draft
wants to merge 36 commits into
base: main
Choose a base branch
from
Draft

fix(typing): Overhaul @overloads in nw.from_native #3125

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
36 commits
Select commit Hold shift + click to select a range
d5a126d
fix(typing): Huuuuuuuge progress
dangotbanned Sep 12, 2025
39bb6b3
test(typing): Add `pa.ChunkedArray` into the mix
dangotbanned Sep 12, 2025
c7cf5ae
test: Rename variables
dangotbanned Sep 12, 2025
6f56dc0
test: Cover w/ `pass_through=True` as well
dangotbanned Sep 12, 2025
d72ce2d
test: Use `assert_type` and runtime checks
dangotbanned Sep 12, 2025
94029ed
test: rename test
dangotbanned Sep 12, 2025
0358cc9
test: Insane `LazyFrame` coverage
dangotbanned Sep 12, 2025
43ceb3f
revert: wont be using that one
dangotbanned Sep 12, 2025
ceaf9e7
refactor(typing): Factor out another pair
dangotbanned Sep 12, 2025
d988a24
refactor(typing): Add single `pass_through` case
dangotbanned Sep 12, 2025
cc5588b
revert: never used them
dangotbanned Sep 12, 2025
1cab6d0
refactor: Invert `NotRequired` -> `Required`
dangotbanned Sep 12, 2025
ff6258d
refactor: Rename all `TypedDict`s
dangotbanned Sep 12, 2025
a0095b1
style: Remove whitespace
dangotbanned Sep 12, 2025
06d74a3
Merge remote-tracking branch 'upstream/main' into from-native-overloads
dangotbanned Sep 12, 2025
c02ab37
Merge branch 'main' into from-native-overloads
dangotbanned Sep 13, 2025
5b6ca5f
Merge branch 'main' into from-native-overloads
dangotbanned Oct 6, 2025
4142e43
Merge remote-tracking branch 'upstream/main' into from-native-overloads
dangotbanned Oct 20, 2025
2e7a0ae
chore: fix merge conflicts
dangotbanned Oct 20, 2025
435beb4
refactor: Move new `TypedDict`s to `_translate.py`
dangotbanned Oct 20, 2025
0605936
fix(typing): More gracefully handle narwhals in overloads
dangotbanned Oct 20, 2025
67bb52f
revert: Remove now-unused `kwds`
dangotbanned Oct 20, 2025
fd7ab16
refactor(typing): Update `v2`
dangotbanned Oct 20, 2025
9a8c025
refactor(DRAFT): Start shrinking `v1`
dangotbanned Oct 20, 2025
c61af03
refactor(typing): Add `AllowAny` variants
dangotbanned Oct 20, 2025
252bcdb
refactor(typing): Add `AllowSeries` variants
dangotbanned Oct 20, 2025
c887618
refactor(typing): Add `OnlySeries` variants
dangotbanned Oct 20, 2025
99527d5
refactor(typing): Add `ExcludeSeries` variants
dangotbanned Oct 20, 2025
1dc56f0
refactor(typing): Add `AllowLazy` variants
dangotbanned Oct 20, 2025
10e7ddf
refactor(typing): Add `OnlyEagerOrInterchange` variants
dangotbanned Oct 20, 2025
1c498aa
chore: remove notes
dangotbanned Oct 20, 2025
4e37df6
chore: Fix those imports for `v1` too
dangotbanned Oct 20, 2025
f372696
oop missed a spot
dangotbanned Oct 20, 2025
459598f
chore: remove some comments
dangotbanned Oct 20, 2025
58590e4
Merge remote-tracking branch 'upstream/main' into from-native-overloads
dangotbanned Oct 20, 2025
bd90c1f
Merge branch 'main' into from-native-overloads
dangotbanned Oct 23, 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
21 changes: 2 additions & 19 deletions narwhals/stable/v2/typing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Any, Protocol, TypeVar, Union
from typing import TYPE_CHECKING, Any, TypeVar, Union

if TYPE_CHECKING:
import sys
from collections.abc import Iterable, Sized

from narwhals.stable.v2 import DataFrame, LazyFrame

Expand All @@ -14,23 +13,7 @@
from typing_extensions import TypeAlias

from narwhals.stable.v2 import Expr, Series

# All dataframes supported by Narwhals have a
# `columns` property. Their similarities don't extend
# _that_ much further unfortunately...
class NativeFrame(Protocol):
@property
def columns(self) -> Any: ...

def join(self, *args: Any, **kwargs: Any) -> Any: ...

class NativeDataFrame(Sized, NativeFrame, Protocol): ...

class NativeLazyFrame(NativeFrame, Protocol):
def explain(self, *args: Any, **kwargs: Any) -> Any: ...

class NativeSeries(Sized, Iterable[Any], Protocol):
def filter(self, *args: Any, **kwargs: Any) -> Any: ...
from narwhals.typing import NativeDataFrame, NativeLazyFrame, NativeSeries


IntoExpr: TypeAlias = Union["Expr", str, "Series[Any]"]
Expand Down
173 changes: 55 additions & 118 deletions narwhals/translate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import datetime as dt
from decimal import Decimal
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, overload
from typing import TYPE_CHECKING, Any, Callable, Literal, TypedDict, TypeVar, overload

from narwhals._constants import EPOCH, MS_PER_SECOND
from narwhals._namespace import (
Expand All @@ -30,6 +30,8 @@
)

if TYPE_CHECKING:
from typing_extensions import Required, Unpack

from narwhals.dataframe import DataFrame, LazyFrame
from narwhals.series import Series
from narwhals.typing import (
Expand Down Expand Up @@ -99,150 +101,87 @@ def to_native(
return narwhals_object


@overload
def from_native(native_object: SeriesT, **kwds: Any) -> SeriesT: ...
class ExcludeSeries(TypedDict, total=False):
pass_through: bool
eager_only: bool
series_only: Literal[False]
allow_series: Literal[False] | None


@overload
def from_native(native_object: DataFrameT, **kwds: Any) -> DataFrameT: ...
class AllowSeries(TypedDict, total=False):
pass_through: bool
eager_only: bool
series_only: Literal[False]
allow_series: Required[Literal[True]]


@overload
def from_native(native_object: LazyFrameT, **kwds: Any) -> LazyFrameT: ...
class OnlySeries(TypedDict, total=False):
pass_through: bool
eager_only: bool
series_only: Required[Literal[True]]
allow_series: bool | None


@overload
def from_native(
native_object: IntoDataFrameT | IntoSeriesT,
*,
pass_through: Literal[True],
eager_only: Literal[True],
series_only: Literal[False] = ...,
allow_series: Literal[True],
) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ...
class OnlyEager(TypedDict, total=False):
pass_through: bool
eager_only: Required[Literal[True]]
series_only: bool
allow_series: bool | None
Copy link
Member Author

Choose a reason for hiding this comment

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

Issue 2

There's too many combinations we need to account for in @overloads.

  • 3x different, true native_objects
  • 3x narwhals objects pretending to be native
  • 4x flags
    • They all have a default, so needs ...
    • allow_series: bool | None just adds to it

So instead of that:

  • I've narrowed it down to 1-2 keywords that define unique behavior
  • Given that a name
  • Made a TypedDict
  • Then relied on Unpack and Required to merge a whole bunch of them



@overload
def from_native(
native_object: IntoDataFrameT,
*,
pass_through: Literal[True],
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> DataFrame[IntoDataFrameT]: ...
class AllowLazy(TypedDict, total=False):
pass_through: bool
eager_only: Literal[False]
series_only: Literal[False]
allow_series: bool | None


@overload
def from_native(
native_object: T,
*,
pass_through: Literal[True],
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> T: ...
class AllowAny(TypedDict, total=False):
pass_through: bool
eager_only: Literal[False]
series_only: Literal[False]
allow_series: Required[Literal[True]]


@overload
def from_native(
native_object: IntoDataFrameT,
*,
pass_through: Literal[True],
eager_only: Literal[True],
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> DataFrame[IntoDataFrameT]: ...
class PassThroughUnknown(TypedDict, total=False):
pass_through: Required[Literal[True]]
eager_only: bool
series_only: bool
allow_series: bool | None


@overload
def from_native(
native_object: T,
*,
pass_through: Literal[True],
eager_only: Literal[True],
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> T: ...


def from_native(native_object: SeriesT, **kwds: Any) -> SeriesT: ...
@overload
def from_native(native_object: DataFrameT, **kwds: Any) -> DataFrameT: ...
@overload
def from_native(native_object: LazyFrameT, **kwds: Any) -> LazyFrameT: ...
@overload
def from_native(
native_object: IntoDataFrameT | IntoLazyFrameT | IntoSeriesT,
*,
pass_through: Literal[True],
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: Literal[True],
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT]: ...


native_object: IntoDataFrameT, **kwds: Unpack[ExcludeSeries]
) -> DataFrame[IntoDataFrameT]: ...
@overload
def from_native(
native_object: IntoSeriesT,
*,
pass_through: Literal[True],
eager_only: Literal[False] = ...,
series_only: Literal[True],
allow_series: None = ...,
native_object: IntoSeriesT, **kwds: Unpack[OnlySeries]
) -> Series[IntoSeriesT]: ...


@overload
def from_native(
native_object: IntoLazyFrameT,
*,
pass_through: Literal[False] = ...,
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> LazyFrame[IntoLazyFrameT]: ...


native_object: IntoSeriesT, **kwds: Unpack[AllowSeries]
) -> Series[IntoSeriesT]: ...
@overload
def from_native(
native_object: IntoDataFrameT,
*,
pass_through: Literal[False] = ...,
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> DataFrame[IntoDataFrameT]: ...


native_object: IntoLazyFrameT, **kwds: Unpack[AllowLazy]
) -> LazyFrame[IntoLazyFrameT]: ...
@overload
def from_native(
native_object: IntoDataFrameT,
*,
pass_through: Literal[False] = ...,
eager_only: Literal[True],
series_only: Literal[False] = ...,
allow_series: None = ...,
) -> DataFrame[IntoDataFrameT]: ...


native_object: IntoDataFrameT | IntoSeriesT, **kwds: Unpack[AllowSeries]
) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ...
@overload
def from_native(
native_object: IntoFrame | IntoSeries,
*,
pass_through: Literal[False] = ...,
eager_only: Literal[False] = ...,
series_only: Literal[False] = ...,
allow_series: Literal[True],
) -> DataFrame[Any] | LazyFrame[Any] | Series[Any]: ...


native_object: IntoDataFrameT | IntoLazyFrameT | IntoSeriesT, **kwds: Unpack[AllowAny]
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT]: ...
@overload
def from_native(
native_object: IntoSeriesT,
*,
pass_through: Literal[False] = ...,
eager_only: Literal[False] = ...,
series_only: Literal[True],
allow_series: None = ...,
) -> Series[IntoSeriesT]: ...


def from_native(native_object: T, **kwds: Unpack[PassThroughUnknown]) -> T: ...
# All params passed in as variables
@overload
def from_native(
Expand All @@ -253,8 +192,6 @@ def from_native(
series_only: bool,
allow_series: bool | None,
) -> Any: ...


def from_native( # noqa: D417
native_object: IntoLazyFrameT
| IntoDataFrameT
Expand Down
8 changes: 7 additions & 1 deletion narwhals/typing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,19 @@ def columns(self) -> Any: ...

def join(self, *args: Any, **kwargs: Any) -> Any: ...

class NativeDataFrame(Sized, NativeFrame, Protocol): ...
class NativeDataFrame(Sized, NativeFrame, Protocol):
def drop(self, *args: Any, **kwargs: Any) -> Any: ...

class NativeLazyFrame(NativeFrame, Protocol):
def explain(self, *args: Any, **kwargs: Any) -> Any: ...

# Needs to have something `NativeDataFrame` doesn't?
class NativeSeries(Sized, Iterable[Any], Protocol):
def filter(self, *args: Any, **kwargs: Any) -> Any: ...
# `pd.DataFrame` has this - the others don't
def value_counts(self, *args: Any, **kwargs: Any) -> Any: ...
# `pl.DataFrame` has this - the others don't
def unique(self, *args: Any, **kwargs: Any) -> Any: ...
Copy link
Member Author

@dangotbanned dangotbanned Sep 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Issue 1

NativeDataFrame and NativeSeries could overlap.

Although we don't define NativeDataFrame.filter - the classes that match it do:

That was leading to the strange @overload order and parameters changing the return type.

But now if we started with NativeSeries, then both series_only and allow_series will preserve that detail:

Prove it bro

narwhals/tests/translate/from_native_test.py

Lines 621 to 723 in a0095b1

def test_from_native_series_exhaustive() -> None: # noqa: PLR0914, PLR0915
pytest.importorskip("polars")
pytest.importorskip("pandas")
pytest.importorskip("pyarrow")
pytest.importorskip("typing_extensions")
import pandas as pd
import polars as pl
import pyarrow as pa
from typing_extensions import assert_type
pl_ser = pl.Series([1, 2, 3])
pd_ser = cast("pd.Series[Any]", pd.Series([1, 2, 3]))
pa_ser = cast("pa.ChunkedArray[Any]", pa.chunked_array([pa.array([1])])) # type: ignore[redundant-cast]
pl_1 = nw.from_native(pl_ser, series_only=True)
pl_2 = nw.from_native(pl_ser, allow_series=True)
pl_3 = nw.from_native(pl_ser, eager_only=True, series_only=True)
pl_4 = nw.from_native(pl_ser, eager_only=True, series_only=True, allow_series=True)
pl_5 = nw.from_native(pl_ser, eager_only=True, allow_series=True)
pl_6 = nw.from_native(pl_ser, series_only=True, allow_series=True)
pl_7 = nw.from_native(pl_ser, series_only=True, pass_through=True)
pl_8 = nw.from_native(pl_ser, allow_series=True, pass_through=True)
pl_9 = nw.from_native(pl_ser, eager_only=True, series_only=True, pass_through=True)
pl_10 = nw.from_native(
pl_ser, eager_only=True, series_only=True, allow_series=True, pass_through=True
)
pl_11 = nw.from_native(pl_ser, eager_only=True, allow_series=True, pass_through=True)
pl_12 = nw.from_native(pl_ser, series_only=True, allow_series=True, pass_through=True)
pls = pl_1, pl_2, pl_3, pl_4, pl_5, pl_6, pl_7, pl_8, pl_9, pl_10, pl_11, pl_12
assert_type(pl_1, nw.Series[pl.Series])
assert_type(pl_2, nw.Series[pl.Series])
assert_type(pl_3, nw.Series[pl.Series])
assert_type(pl_4, nw.Series[pl.Series])
assert_type(pl_5, nw.Series[pl.Series])
assert_type(pl_6, nw.Series[pl.Series])
assert_type(pl_7, nw.Series[pl.Series])
assert_type(pl_8, nw.Series[pl.Series])
assert_type(pl_9, nw.Series[pl.Series])
assert_type(pl_10, nw.Series[pl.Series])
assert_type(pl_11, nw.Series[pl.Series])
assert_type(pl_12, nw.Series[pl.Series])
pd_1 = nw.from_native(pd_ser, series_only=True)
pd_2 = nw.from_native(pd_ser, allow_series=True)
pd_3 = nw.from_native(pd_ser, eager_only=True, series_only=True)
pd_4 = nw.from_native(pd_ser, eager_only=True, series_only=True, allow_series=True)
pd_5 = nw.from_native(pd_ser, eager_only=True, allow_series=True)
pd_6 = nw.from_native(pd_ser, series_only=True, allow_series=True)
pd_7 = nw.from_native(pd_ser, series_only=True, pass_through=True)
pd_8 = nw.from_native(pd_ser, allow_series=True, pass_through=True)
pd_9 = nw.from_native(pd_ser, eager_only=True, series_only=True, pass_through=True)
pd_10 = nw.from_native(
pd_ser, eager_only=True, series_only=True, allow_series=True, pass_through=True
)
pd_11 = nw.from_native(pd_ser, eager_only=True, allow_series=True, pass_through=True)
pd_12 = nw.from_native(pd_ser, series_only=True, allow_series=True, pass_through=True)
pds = pd_1, pd_2, pd_3, pd_4, pd_5, pd_6, pd_7, pd_8, pd_9, pd_10, pd_11, pd_12
assert_type(pd_1, nw.Series["pd.Series[Any]"])
assert_type(pd_2, nw.Series["pd.Series[Any]"])
assert_type(pd_3, nw.Series["pd.Series[Any]"])
assert_type(pd_4, nw.Series["pd.Series[Any]"])
assert_type(pd_5, nw.Series["pd.Series[Any]"])
assert_type(pd_6, nw.Series["pd.Series[Any]"])
assert_type(pd_7, nw.Series["pd.Series[Any]"])
assert_type(pd_8, nw.Series["pd.Series[Any]"])
assert_type(pd_9, nw.Series["pd.Series[Any]"])
assert_type(pd_10, nw.Series["pd.Series[Any]"])
assert_type(pd_11, nw.Series["pd.Series[Any]"])
assert_type(pd_12, nw.Series["pd.Series[Any]"])
pa_1 = nw.from_native(pa_ser, series_only=True)
pa_2 = nw.from_native(pa_ser, allow_series=True)
pa_3 = nw.from_native(pa_ser, eager_only=True, series_only=True)
pa_4 = nw.from_native(pa_ser, eager_only=True, series_only=True, allow_series=True)
pa_5 = nw.from_native(pa_ser, eager_only=True, allow_series=True)
pa_6 = nw.from_native(pa_ser, series_only=True, allow_series=True)
pa_7 = nw.from_native(pa_ser, series_only=True, pass_through=True)
pa_8 = nw.from_native(pa_ser, allow_series=True, pass_through=True)
pa_9 = nw.from_native(pa_ser, eager_only=True, series_only=True, pass_through=True)
pa_10 = nw.from_native(
pa_ser, eager_only=True, series_only=True, allow_series=True, pass_through=True
)
pa_11 = nw.from_native(pa_ser, eager_only=True, allow_series=True, pass_through=True)
pa_12 = nw.from_native(pa_ser, series_only=True, allow_series=True, pass_through=True)
pas = pa_1, pa_2, pa_3, pa_4, pa_5, pa_6, pa_7, pa_8, pa_9, pa_10, pa_11, pa_12
assert_type(pa_1, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_2, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_3, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_4, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_5, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_6, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_7, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_8, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_9, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_10, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_11, nw.Series["pa.ChunkedArray[Any]"])
assert_type(pa_12, nw.Series["pa.ChunkedArray[Any]"])
for series in chain(pls, pds, pas):
assert isinstance(series, nw.Series)


class SupportsNativeNamespace(Protocol):
def __native_namespace__(self) -> ModuleType: ...
Expand Down
Loading
Loading