fix(typing): Overhaul @overloads in nw.from_native
#3125
Conversation
Will close #1928
Absolutely not doing this with another class
wayyy more common in this case
The `TypedDict`s are clear enough
narwhals/typing.py
Outdated
| 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: ... |
There was a problem hiding this comment.
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:
narwhals/translate.py
Outdated
| class OnlyEager(TypedDict, total=False): | ||
| pass_through: bool | ||
| eager_only: Required[Literal[True]] | ||
| series_only: bool | ||
| allow_series: bool | None |
There was a problem hiding this comment.
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 | Nonejust adds to it
- They all have a default, so needs
So instead of that:
dangotbanned
changed the title
from_native@overloads in nw.from_native
|
@MarcoGorelli this one isn't ready for a proper review yet, but I was curious about your thoughts on the idea generally? I did some testing locally and I think depending on how these issues have progressed - docs might be one more area for experimenting |
tbh I'm still not too keen on matching on the native object, i kind of feel like |
That isn't possible due to the issue from: Which is also why Possibly these sections of the spec may be helpful?
These are the issues that come from having such a heavily polymorphic function I'm afraid 😢 |
`v1`, `v2` can extend and reuse these defs
Need `OnlySeries` to finish the other replacements
ooooh I think we're almost ready
I've updated the PR description with a table showing this in more detail
| return super().__getitem__(item) | ||
|
|
||
| def get_column(self, name: str) -> Series: | ||
| def get_column(self, name: str) -> Series[Any]: |
There was a problem hiding this comment.
v1.Series only had a default TypeVar to avoid downstream noise from making it generic
That shouldn't have carried over into v2
| def test_from_native_lazyframe_exhaustive() -> None: # noqa: PLR0914, PLR0915 | ||
| pytest.importorskip("polars") | ||
| pytest.importorskip("typing_extensions") | ||
|
|
||
| import polars as pl | ||
| from typing_extensions import assert_type | ||
|
|
||
| pl_ldf = pl.LazyFrame(data) | ||
|
|
||
| pl_1 = nw.from_native(pl_ldf) |
There was a problem hiding this comment.
Just to get ahead of the obvious question:
Do we need all of this?
The answer is complicated.
IMO, the real source of all the typing issues with from_native has been that we test very little in terms of typing for such a complex function.
I think that the primary entry point for narwhals should require much more care to be sure we don't regress.
That's my opinion, but I wouldn't feel so confident in making such big changes to the typing without this in place.
I'm guessing Marco is familiar with this from numpy?
- https://github.com/numpy/numpy/blob/1aa31099403cbf5ba17dab857666986d57749834/numpy/typing/tests/data/reveal/array_constructors.pyi
- https://github.com/numpy/numpy/blob/1aa31099403cbf5ba17dab857666986d57749834/numpy/typing/tests/data/reveal/ma.pyi
Doing this in .pyi files rather than (mostly) in a TYPE_CHECKING block is another option 🙂
There was a problem hiding this comment.
This issue won't be resolved by this PR, but it does highlight the same gap
We have no typing coverage for the example on the README.md?
import narwhals as nw
from narwhals.typing import IntoFrameT
def agnostic_function(
df_native: IntoFrameT,
date_column: str,
price_column: str,
) -> IntoFrameT:
return (
nw.from_native(df_native)
.group_by(nw.col(date_column).dt.truncate("1mo"))
.agg(nw.col(price_column).mean())
.sort(date_column)
.to_native()
)mypy complains about this already 😳
narwhals\/.py:15: error: Value of type variable "IntoLazyFrameT" of "from_native" cannot be "IntoFrameT" [type-var]
nw.from_native(df_native)
^~~~~~~~~~~~~~~~~~~~~~~~~
Found 1 error in 1 file (checked 1 source file)
|
Thanks @MarcoGorelli
I'm exhausted just looking at them |
|
aw i'm sorry for this api 😢 |
Don't be silly, I'm dunking on myself here 😄 |
Closes #1928
What type of PR is this? (check all applicable)
Related issues
DataFrame,LazyFrame#2944Staying in draftLeft in draft until refactor(typing): Add_native.py#3086 was mergedIf you have comments or can explain your changes, please do so below
Note
The positive diff is from new quite exhaustive tests
narwhals.translatehas shrunk a lot 🥳Show tasks
Native{DataFrame,Series}nw._nativeNativeSeriesfiltermight not be needed anymoreTypedDicts intonw._translateTypedDicts forstablev1(thread)v1AllowAnyvariantsAllowSeriesvariantsOnlySeriesvariantsExcludeSeriesvariantsAllowLazyvariantsOnlyEagerOrInterchangevariantsv2mainfrom_native()@overloads 1-3 to useTypedDictsmkdocspresents the new typingTypedDicts?Unpackwork out-of-the-box?@overloads useful/noisy?Previews
I've collected a bunch of things here that may be missed from the diff alone.
Be sure to check out the
<details>😉fix(typing): More gracefully handle narwhals in overloads
Anywas previously being used on the first few overloads to solve a recursion issue.But that had the downside of swallowing completions until you matched an overload with
native_object:Show before
But now, we have something more useful
Show after
Changes to
Native{Series,DataFrame}As mentioned in (#3125 (comment)), these protocols needed to be
adapted to make them disjoint across
pandas,polars,pyarrow:dropfilteruniquevalue_countspd.Seriespl.Seriespa.ChunkedArraypd.DataFramepl.DataFramepa.TableMkDocs Comparison
These changes shrink the API docs for
from_nativequite a lot.However, if the specific arguments combinations are wanted here - then an options is adding the
TypedDicts somewhere and enablingsignature_crossrefs.IIRC, that feature used to be insiders-only but looks to be available since https://squidfunk.github.io/mkdocs-material/blog/2025/11/11/insiders-now-free-for-everyone/
Show
mainShow (
636f787)Note
I'm fine with less detail here, but
signature_crossrefsis an option if this is blocking 🙂