fix(typing): Narrow `TypeVar`(s) used in `(Data|Lazy)Frame` #2356

dangotbanned · 2025-04-07T13:10:31Z

Will close #2344, #2345

What type of PR is this? (check all applicable)

Related issues

Closes Narrow TypeVar used in DataFrame #2344
Closes Narrow TypeVar used in LazyFrame #2345
Related [Enh]: Add pivot to Arrow #2179

Checklist

Code follows style guide (ruff)
Tests added
Documented the changes

If you have comments or can explain your changes, please do so below

Will close #2344

- Revealed two new `[overload-cannot-match]` from `mypy` - I agreed with that and removed the conflict sources Will close #2345

narwhals/translate.py

Revealed quite a few other issues

+ fix related quirks

- I think this whole test needs rewriting - We shouldn't be depending on the internals like this

https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.to_numpy.html

Spent waaaaaay too long trying to get this working

dangotbanned · 2025-04-07T20:34:08Z

tests/translate/from_native_test.py

+def test_dataframe_recursive_v1() -> None:
+    pytest.importorskip("polars")
+    import polars as pl
+
+    pl_frame = pl.DataFrame({"a": [1, 2, 3]})
+    nw_frame = nw.from_native(pl_frame)
+    with pytest.raises(AssertionError):
+        nw.DataFrame(nw_frame, level="full")
+
+    nw_frame_early_return = nw.from_native(nw_frame)
+
+    if TYPE_CHECKING:
+        assert_type(pl_frame, pl.DataFrame)
+        # TODO @dangotbanned: Fix without breaking something else (1)
+        assert_type(nw_frame, nw.DataFrame[pl.DataFrame])
+
+        nw_frame_depth_2 = nw.DataFrame(nw_frame, level="full")
+        # NOTE: Checking that the type is `DataFrame[Unknown]`
+        assert_type(nw_frame_depth_2, nw.DataFrame)
+        # TODO @dangotbanned: Fix without breaking something else (2)
+        assert_type(nw_frame_early_return, nw.DataFrame[pl.DataFrame])


Spent a few hours trying to get this working

Nothing I've tried worked for both mypy, pyright and didn't break something elsewhere

AFAICT, v1s overloads were always returning a Union for the case I was trying to force into DataFrame.

Gave up on trying to resolve in (c4bed59)

#2356 (comment)

dangotbanned · 2025-04-07T21:08:09Z

tests/translate/from_native_test.py

@@ -1,5 +1,6 @@
 from __future__ import annotations

+# mypy: disallow-any-generics=false, disable-error-code="var-annotated"


Needed to add this to avoid mypy requiring # type: ignores.

If a line has one (regardless of it is has an error code) pyright won't ever emit a diagnostic

Moved into a docstring for a better explainer (6a66779)

@FBruzzesi in relation to (#2304 (comment))

I think there's quite a lot of things to be learned - by following this PR all the way back to the original issue (#2239) 🙂

Thanks for the ping! I just went through the entire PR, my skill issue is that everything makes sense, but most likely I would not be able to (re)produce it 🥲

@FBruzzesi so if you go through the PR (or (https://github.com/narwhals-dev/narwhals/pull/2347/commits)) in order of commits - there's a pattern in all of this.

Is this a workflow?

Identify some typing that doesn't meet your expectation ([Bug]: TypeVar(s) used in nw.(BaseFrame|Series) are recursive #2239)

Make the smallest change you think could fix it (74a03a8)

Write a test that would demonstrate the bug was fixed (ff841d1)

I probably should've started with that, but I was eager I guess 😉

Add some notes when that didn't work (ff841d1#r2029276670)

Add even more notes when all seems lost (37549cc)

Try absolutely anything to fix it (6bd47b4)

I wasn't happy with this initially, but became more convinced when it didn't break anything else

Without the tests in place - I might have not even tried this

Write an essay on why you think that worked (6bd47b4#r2029455047)

Clean things up (50e4d13)

Back to this PR

You can follow through (aa6578e) and Next > to see a similar thing 🙂

Different issues came up, but I'm mostly trying to isolate small things and see what mypy & pyright spits back out when I attempt a fix

I also usually try to get the typing right first before I start making "real" code changes.

It is much harder to spot the source of an issue if you only check the types after making lots of changes

narwhals/typing.py

#2356 (comment)

Expanded on (#2356 (comment))

MarcoGorelli

thanks @dangotbanned ! just got a question

MarcoGorelli · 2025-04-09T14:38:23Z

narwhals/_polars/dataframe.py

-        index: list[str] | None,
-        values: list[str] | None,
+        index: str | Sequence[str] | None,
+        values: str | Sequence[str] | None,


why do these change? is Sequence[str] | None not fine?

thanks @dangotbanned ! just got a question

Thanks @MarcoGorelli, happy to answer 😅

why do these change? is Sequence[str] | None not fine?

the original annotation was overly narrow, since polars accepts much more than a list

https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.pivot.html

https://github.com/pola-rs/polars/blob/py-1.26.0/py-polars/polars/dataframe/frame.py#L8732-L8963

I followed through to pandas to see if that was the source of using list[str], but IIRC it used an alias like IndexLabel.

Sequence[str] | None is equivalent to str | Sequence[str] | None sadly (python/typing#256).
But I prefer being more explicit that it can accept "one or more columns"

if it's str we already make it a list at the narwhals level, can this just be Sequence[str] | None?

@MarcoGorelli sure thing!

Resolved in (d6cb16b)

https://github.com/narwhals-dev/narwhals/pull/2356/files/5dd782522f23ed2aef3554a2aa89fc9903abd094#r2040702116

#2352

dangotbanned · 2025-04-12T19:34:22Z

@MarcoGorelli I've resolved the conflicts from #2380

Are we good to go on this one?

dangotbanned · 2025-04-14T17:49:49Z

narwhals/_compliant/group_by.py

+class DataFrameGroupBy(
+    CompliantGroupBy[CompliantDataFrameT_co, CompliantExprT_contra],
+    Protocol38[CompliantDataFrameT_co, CompliantExprT_contra],
+):
+    def __iter__(self) -> Iterator[tuple[Any, CompliantDataFrameT_co]]: ...


@FBruzzesi I forgot to mention this will probably be helpful in getting the typing working for (#2325).

It means you can use CompliantGroupBy for Polars* and put the unrelated parts here.

I needed to add it to resolve an unrelated issue (1815752)

MarcoGorelli

thanks @dangotbanned !

sorry for the delay

dangotbanned · 2025-04-15T19:41:29Z

thanks @dangotbanned !

sorry for the delay

No worries @MarcoGorelli!

I'm just happy we're now in a good place to work on

API: io functions for v2 #2116

😊

dangotbanned added 3 commits April 7, 2025 13:31

fix(typing): Narrow IntoDataFrame

aa6578e

Will close #2344

fix(typing): Remove DataFrame from IntoFrame

51c5b63

fix(typing): Narrow IntoLazyFrame, IntoFrame

6cbbc74

- Revealed two new `[overload-cannot-match]` from `mypy` - I agreed with that and removed the conflict sources Will close #2345

dangotbanned added fix typing labels Apr 7, 2025

dangotbanned commented Apr 7, 2025

View reviewed changes

narwhals/translate.py Show resolved Hide resolved

dangotbanned commented Apr 7, 2025

View reviewed changes

narwhals/translate.py Show resolved Hide resolved

dangotbanned added 10 commits April 7, 2025 14:49

fix(typing): Annotate DataFrame._compliant_frame

255ab27

Revealed quite a few other issues

chore: Add missing CompliantDataFrame.pivot

ba2f6e1

+ fix related quirks

fix(typing): Ensure __iter__ is available on group_by

1815752

chore(typing): Fix most of DataFrame

07deea2

chore(typing): Ignore interchange [type-var]

3881822

test(typing): Barely fix dodgy spark typing

375fabc

- I think this whole test needs rewriting - We shouldn't be depending on the internals like this

fix: Implement to_numpy to catch args

21e80ef

https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.to_numpy.html

fix(typing): Annotate LazyFrame._compliant_frame

c124985

chore(typing): Ignore and add note for spark_like cast

831a6be

chore(typing): Partial v1 backport

1725f36

Spent waaaaaay too long trying to get this working

dangotbanned commented Apr 7, 2025

View reviewed changes

dangotbanned added 2 commits April 7, 2025 21:53

fix(typing): Just preserve v1 behavior

c4bed59

#2356 (comment)

simplify

6a9fd91

dangotbanned commented Apr 7, 2025

View reviewed changes

narwhals/typing.py Outdated Show resolved Hide resolved

try old Union

ed65ad2

#2356 (comment)

dangotbanned marked this pull request as ready for review April 7, 2025 21:49

dangotbanned added 2 commits April 8, 2025 11:16

Merge remote-tracking branch 'upstream/main' into narrow-type-var-frame

b97149d

docs(typing): Provide more context on what and why

6a66779

Expanded on (#2356 (comment))

dangotbanned mentioned this pull request Apr 8, 2025

feat: Adds private Namespace class #2324

Merged

10 tasks

dangotbanned added 2 commits April 8, 2025 22:51

Merge branch 'main' into narrow-type-var-frame

675329c

Merge branch 'main' into narrow-type-var-frame

5dd7825

dangotbanned requested a review from MarcoGorelli April 9, 2025 09:51

MarcoGorelli reviewed Apr 9, 2025

View reviewed changes

dangotbanned added 9 commits April 9, 2025 16:20

Merge branch 'main' into narrow-type-var-frame

a1c51ff

Merge branch 'main' into narrow-type-var-frame

c2b328b

Merge branch 'main' into narrow-type-var-frame

b66d1bf

Merge branch 'main' into narrow-type-var-frame

ba525f1

Merge branch 'main' into narrow-type-var-frame

c45f1f4

Merge branch 'main' into narrow-type-var-frame

322b00d

chore(typing): Use Sequence[str] in pivot

d6cb16b

https://github.com/narwhals-dev/narwhals/pull/2356/files/5dd782522f23ed2aef3554a2aa89fc9903abd094#r2040702116

refactor(typing): Use PivotAgg

cbd60d9

#2352

Merge remote-tracking branch 'upstream/main' into narrow-type-var-frame

53722e8

Merge branch 'main' into narrow-type-var-frame

e111fc3

dangotbanned requested a review from MarcoGorelli April 13, 2025 12:24

dangotbanned added 2 commits April 13, 2025 21:32

Merge branch 'main' into narrow-type-var-frame

23158d1

Merge branch 'main' into narrow-type-var-frame

bc72941

dangotbanned requested a review from FBruzzesi April 14, 2025 17:45

dangotbanned commented Apr 14, 2025

View reviewed changes

Merge branch 'main' into narrow-type-var-frame

a1fb349

MarcoGorelli approved these changes Apr 15, 2025

View reviewed changes

MarcoGorelli merged commit 24bdf86 into main Apr 15, 2025
28 of 29 checks passed

MarcoGorelli deleted the narrow-type-var-frame branch April 15, 2025 19:33

dangotbanned mentioned this pull request Apr 15, 2025

[Bug]: TypeVar(s) used in nw.(BaseFrame|Series) are recursive #2239

Closed

dangotbanned mentioned this pull request Aug 5, 2025

fix(typing): Avoid overlapping DataFrame, LazyFrame #2944

Merged

10 tasks

dangotbanned mentioned this pull request Sep 15, 2025

feat: explain or similar API for inspecting plans ibis-project/ibis#9276

Open

1 task

		@@ -1,5 +1,6 @@
		from __future__ import annotations

		# mypy: disallow-any-generics=false, disable-error-code="var-annotated"

fix(typing): Narrow TypeVar(s) used in (Data|Lazy)Frame #2356

fix(typing): Narrow TypeVar(s) used in (Data|Lazy)Frame #2356

Uh oh!

Conversation

dangotbanned commented Apr 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What type of PR is this? (check all applicable)

Related issues

Checklist

If you have comments or can explain your changes, please do so below

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned Apr 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Is this a workflow?

Back to this PR

Uh oh!

dangotbanned Apr 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MarcoGorelli left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned commented Apr 12, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MarcoGorelli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dangotbanned commented Apr 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

fix(typing): Narrow `TypeVar`(s) used in `(Data|Lazy)Frame` #2356

fix(typing): Narrow `TypeVar`(s) used in `(Data|Lazy)Frame` #2356

dangotbanned commented Apr 7, 2025 •

edited

Loading

dangotbanned Apr 7, 2025 •

edited

Loading

dangotbanned Apr 8, 2025 •

edited

Loading