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

* fix(typing): Narrow `IntoDataFrame`

Will close #2344

* fix(typing): Remove `DataFrame` from `IntoFrame`

* fix(typing): Narrow `IntoLazyFrame`, `IntoFrame`

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

Will close #2345

* fix(typing): Annotate `DataFrame._compliant_frame`

Revealed quite a few other issues

* chore: Add missing `CompliantDataFrame.pivot`

+ fix related quirks

* fix(typing): Ensure `__iter__` is available on group_by

* chore(typing): Fix most of `DataFrame`

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

* test(typing): Barely fix dodgy spark typing

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

* fix: Implement `to_numpy` to catch args

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

* fix(typing): Annotate `LazyFrame._compliant_frame`

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

* chore(typing): Partial `v1` backport

Spent waaaaaay too long trying to get this working

* fix(typing): Just preserve `v1` behavior

#2356 (comment)

* simplify

* try old `Union`

#2356 (comment)

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

Expanded on (#2356 (comment))

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

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

* refactor(typing): Use `PivotAgg`

#2352

Author: dangotbanned

narwhals/_arrow/dataframe.py

@@ -845,3 +845,5 @@ def unpivot(
         )
         # TODO(Unassigned): Even with promote_options="permissive", pyarrow does not
         # upcast numeric to non-numeric (e.g. string) datatypes
+
+    pivot = not_implemented()

narwhals/_compliant/dataframe.py

@@ -37,12 +37,14 @@
     from typing_extensions import TypeAlias
 
     from narwhals._compliant.group_by import CompliantGroupBy
+    from narwhals._compliant.group_by import DataFrameGroupBy
     from narwhals._translate import IntoArrowTable
     from narwhals.dtypes import DType
     from narwhals.schema import Schema
     from narwhals.typing import AsofJoinStrategy
     from narwhals.typing import JoinStrategy
     from narwhals.typing import LazyUniqueKeepStrategy
+    from narwhals.typing import PivotAgg
     from narwhals.typing import SizeUnit
     from narwhals.typing import UniqueKeepStrategy
     from narwhals.typing import _2DArray
@@ -136,7 +138,7 @@ def gather_every(self, n: int, offset: int) -> Self: ...
     def get_column(self, name: str) -> CompliantSeriesT: ...
     def group_by(
         self, *keys: str, drop_null_keys: bool
-    ) -> CompliantGroupBy[Self, Any]: ...
+    ) -> DataFrameGroupBy[Self, Any]: ...
     def head(self, n: int) -> Self: ...
     def item(self, row: int | None, column: int | str | None) -> Any: ...
     def iter_columns(self) -> Iterator[CompliantSeriesT]: ...
@@ -165,6 +167,16 @@ def join_asof(
         suffix: str,
     ) -> Self: ...
     def lazy(self, *, backend: Implementation | None) -> CompliantLazyFrame[Any, Any]: ...
+    def pivot(
+        self,
+        on: Sequence[str],
+        *,
+        index: Sequence[str] | None,
+        values: Sequence[str] | None,
+        aggregate_function: PivotAgg | None,
+        sort_columns: bool,
+        separator: str,
+    ) -> Self: ...
     def rename(self, mapping: Mapping[str, str]) -> Self: ...
     def row(self, index: int) -> tuple[Any, ...]: ...
     def rows(

narwhals/_compliant/group_by.py

@@ -76,6 +76,13 @@ def __init__(
     def agg(self, *exprs: CompliantExprT_contra) -> CompliantFrameT_co: ...
 
 
+class DataFrameGroupBy(
+    CompliantGroupBy[CompliantDataFrameT_co, CompliantExprT_contra],
+    Protocol38[CompliantDataFrameT_co, CompliantExprT_contra],
+):
+    def __iter__(self) -> Iterator[tuple[Any, CompliantDataFrameT_co]]: ...
+
+
 class DepthTrackingGroupBy(
     CompliantGroupBy[CompliantFrameT_co, DepthTrackingExprT_contra],
     Protocol38[CompliantFrameT_co, DepthTrackingExprT_contra, NativeAggregationT_co],
@@ -132,9 +139,9 @@ def _leaf_name(cls, expr: DepthTrackingExprAny, /) -> NarwhalsAggregation | Any:
 
 class EagerGroupBy(
     DepthTrackingGroupBy[CompliantDataFrameT_co, EagerExprT_contra, str],
+    DataFrameGroupBy[CompliantDataFrameT_co, EagerExprT_contra],
     Protocol38[CompliantDataFrameT_co, EagerExprT_contra],
-):
-    def __iter__(self) -> Iterator[tuple[Any, CompliantDataFrameT_co]]: ...
+): ...
 
 
 class LazyGroupBy(

narwhals/_pandas_like/dataframe.py

@@ -65,6 +65,7 @@
     from narwhals.typing import CompliantLazyFrame
     from narwhals.typing import DTypeBackend
     from narwhals.typing import JoinStrategy
+    from narwhals.typing import PivotAgg
     from narwhals.typing import SizeUnit
     from narwhals.typing import UniqueKeepStrategy
     from narwhals.typing import _1DArray
@@ -1017,12 +1018,12 @@ def gather_every(self: Self, n: int, offset: int) -> Self:
         return self._with_native(self.native.iloc[offset::n], validate_column_names=False)
 
     def pivot(
-        self: Self,
-        on: list[str],
+        self,
+        on: Sequence[str],
         *,
-        index: list[str] | None,
-        values: list[str] | None,
-        aggregate_function: Any | None,
+        index: Sequence[str] | None,
+        values: Sequence[str] | None,
+        aggregate_function: PivotAgg | None,
         sort_columns: bool,
         separator: str,
     ) -> Self:

narwhals/_pandas_like/utils.py

@@ -5,6 +5,7 @@
 from contextlib import suppress
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Sequence
 from typing import Sized
 from typing import TypeVar
 from typing import cast
@@ -652,9 +653,9 @@ def select_columns_by_name(
 
 def pivot_table(
     df: PandasLikeDataFrame,
-    values: list[str],
-    index: list[str],
-    columns: list[str],
+    values: Sequence[str],
+    index: Sequence[str],
+    columns: Sequence[str],
     aggregate_function: str | None,
 ) -> Any:
     dtypes = import_dtypes_module(df._version)

narwhals/_polars/dataframe.py

@@ -44,6 +44,7 @@
     from narwhals.typing import CompliantDataFrame
     from narwhals.typing import CompliantLazyFrame
     from narwhals.typing import JoinStrategy
+    from narwhals.typing import PivotAgg
     from narwhals.typing import _2DArray
     from narwhals.utils import Version
     from narwhals.utils import _FullContext
@@ -77,7 +78,6 @@ class PolarsDataFrame:
     select: Method[Self]
     sort: Method[Self]
     to_arrow: Method[pa.Table]
-    to_numpy: Method[_2DArray]
     to_pandas: Method[pd.DataFrame]
     unique: Method[Self]
     with_columns: Method[Self]
@@ -232,6 +232,9 @@ def __array__(
             return self.native.__array__(dtype)
         return self.native.__array__(dtype)
 
+    def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _2DArray:
+        return self.native.to_numpy()
+
     def collect_schema(self: Self) -> dict[str, DType]:
         if self._backend_version < (1,):
             return {
@@ -412,15 +415,12 @@ def unpivot(
         )
 
     def pivot(
-        self: Self,
-        on: list[str],
+        self,
+        on: Sequence[str],
         *,
-        index: list[str] | None,
-        values: list[str] | None,
-        aggregate_function: Literal[
-            "min", "max", "first", "last", "sum", "mean", "median", "len"
-        ]
-        | None,
+        index: Sequence[str] | None,
+        values: Sequence[str] | None,
+        aggregate_function: PivotAgg | None,
         sort_columns: bool,
         separator: str,
     ) -> Self: