chore(typing): make some more parts generic

Author: dangotbanned

narwhals/_spark_like/dataframe.py

@@ -257,7 +257,7 @@ def simple_select(self: Self, *column_names: str) -> Self:
 
     def aggregate(
         self: Self,
-        *exprs: SparkLikeExpr,
+        *exprs: SparkLikeExpr[FrameT],
     ) -> Self:
         new_columns = evaluate_exprs(self, *exprs)
 
@@ -266,7 +266,7 @@ def aggregate(
 
     def select(
         self: Self,
-        *exprs: SparkLikeExpr,
+        *exprs: SparkLikeExpr[FrameT],
     ) -> Self:
         new_columns = evaluate_exprs(self, *exprs)
 
@@ -279,11 +279,11 @@ def select(
         new_columns_list = [col.alias(col_name) for (col_name, col) in new_columns]
         return self._from_native_frame(self._native_frame.select(*new_columns_list))
 
-    def with_columns(self: Self, *exprs: SparkLikeExpr) -> Self:
+    def with_columns(self: Self, *exprs: SparkLikeExpr[FrameT]) -> Self:
         new_columns = evaluate_exprs(self, *exprs)
         return self._from_native_frame(self._native_frame.withColumns(dict(new_columns)))
 
-    def filter(self: Self, predicate: SparkLikeExpr) -> Self:
+    def filter(self: Self, predicate: SparkLikeExpr[FrameT]) -> Self:
         # `[0]` is safe as the predicate's expression only returns a single column
         condition = predicate._call(self)[0]
         spark_df = self._native_frame.where(condition)

narwhals/_spark_like/expr.py

@@ -28,22 +28,23 @@
     from pyspark.sql import Window
     from typing_extensions import Self
 
+    from narwhals._spark_like.dataframe import FrameT
     from narwhals._spark_like.dataframe import SparkLikeLazyFrame
     from narwhals._spark_like.namespace import SparkLikeNamespace
     from narwhals._spark_like.typing import WindowFunction
     from narwhals.dtypes import DType
     from narwhals.utils import Version
 
 
-class SparkLikeExpr(CompliantExpr["SparkLikeLazyFrame", "Column"]):  # type: ignore[type-var] # (#2044)
+class SparkLikeExpr(CompliantExpr["SparkLikeLazyFrame[FrameT]", "Column"]):  # type: ignore[type-var] # (#2044)
     _depth = 0  # Unused, just for compatibility with CompliantExpr
 
     def __init__(
         self: Self,
-        call: Callable[[SparkLikeLazyFrame], Sequence[Column]],
+        call: Callable[[SparkLikeLazyFrame[FrameT]], Sequence[Column]],
         *,
         function_name: str,
-        evaluate_output_names: Callable[[SparkLikeLazyFrame], Sequence[str]],
+        evaluate_output_names: Callable[[SparkLikeLazyFrame[FrameT]], Sequence[str]],
         alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
         backend_version: tuple[int, ...],
         version: Version,
@@ -58,11 +59,11 @@ def __init__(
         self._implementation = implementation
         self._window_function: WindowFunction | None = None
 
-    def __call__(self: Self, df: SparkLikeLazyFrame) -> Sequence[Column]:
+    def __call__(self: Self, df: SparkLikeLazyFrame[FrameT]) -> Sequence[Column]:
         return self._call(df)
 
     def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self:
-        def func(df: SparkLikeLazyFrame) -> Sequence[Column]:
+        def func(df: SparkLikeLazyFrame[FrameT]) -> Sequence[Column]:
             if kind is ExprKind.AGGREGATION:
                 return [
                     result.over(df._Window().partitionBy(df._F.lit(1)))
@@ -144,15 +145,15 @@ def __narwhals_namespace__(self: Self) -> SparkLikeNamespace:  # pragma: no cove
     @classmethod
     def from_column_names(
         cls: type[Self],
-        evaluate_column_names: Callable[[SparkLikeLazyFrame], Sequence[str]],
+        evaluate_column_names: Callable[[SparkLikeLazyFrame[FrameT]], Sequence[str]],
         /,
         *,
         function_name: str,
         implementation: Implementation,
         backend_version: tuple[int, ...],
         version: Version,
     ) -> Self:
-        def func(df: SparkLikeLazyFrame) -> list[Column]:
+        def func(df: SparkLikeLazyFrame[FrameT]) -> list[Column]:
             return [df._F.col(col_name) for col_name in evaluate_column_names(df)]
 
         return cls(
@@ -173,7 +174,7 @@ def from_column_indices(
         version: Version,
         implementation: Implementation,
     ) -> Self:
-        def func(df: SparkLikeLazyFrame) -> list[Column]:
+        def func(df: SparkLikeLazyFrame[FrameT]) -> list[Column]:
             columns = df.columns
             return [df._F.col(columns[i]) for i in column_indices]
 
@@ -193,7 +194,7 @@ def _from_call(
         expr_name: str,
         **expressifiable_args: Self | Any,
     ) -> Self:
-        def func(df: SparkLikeLazyFrame) -> list[Column]:
+        def func(df: SparkLikeLazyFrame[FrameT]) -> list[Column]:
             native_series_list = self._call(df)
             other_native_series = {
                 key: maybe_evaluate_expr(df, value)
@@ -230,104 +231,104 @@ def _with_window_function(
         result._window_function = window_function
         return result
 
-    def __eq__(self: Self, other: SparkLikeExpr) -> Self:  # type: ignore[override]
+    def __eq__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:  # type: ignore[override]
         return self._from_call(
             lambda _input, other: _input.__eq__(other), "__eq__", other=other
         )
 
-    def __ne__(self: Self, other: SparkLikeExpr) -> Self:  # type: ignore[override]
+    def __ne__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:  # type: ignore[override]
         return self._from_call(
             lambda _input, other: _input.__ne__(other), "__ne__", other=other
         )
 
-    def __add__(self: Self, other: SparkLikeExpr) -> Self:
+    def __add__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__add__(other), "__add__", other=other
         )
 
-    def __sub__(self: Self, other: SparkLikeExpr) -> Self:
+    def __sub__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__sub__(other), "__sub__", other=other
         )
 
-    def __rsub__(self: Self, other: SparkLikeExpr) -> Self:
+    def __rsub__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: other.__sub__(_input), "__rsub__", other=other
         ).alias("literal")
 
-    def __mul__(self: Self, other: SparkLikeExpr) -> Self:
+    def __mul__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__mul__(other), "__mul__", other=other
         )
 
-    def __truediv__(self: Self, other: SparkLikeExpr) -> Self:
+    def __truediv__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__truediv__(other), "__truediv__", other=other
         )
 
-    def __rtruediv__(self: Self, other: SparkLikeExpr) -> Self:
+    def __rtruediv__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: other.__truediv__(_input), "__rtruediv__", other=other
         ).alias("literal")
 
-    def __floordiv__(self: Self, other: SparkLikeExpr) -> Self:
+    def __floordiv__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         def _floordiv(_input: Column, other: Column) -> Column:
             return self._F.floor(_input / other)
 
         return self._from_call(_floordiv, "__floordiv__", other=other)
 
-    def __rfloordiv__(self: Self, other: SparkLikeExpr) -> Self:
+    def __rfloordiv__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         def _rfloordiv(_input: Column, other: Column) -> Column:
             return self._F.floor(other / _input)
 
         return self._from_call(_rfloordiv, "__rfloordiv__", other=other).alias("literal")
 
-    def __pow__(self: Self, other: SparkLikeExpr) -> Self:
+    def __pow__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__pow__(other), "__pow__", other=other
         )
 
-    def __rpow__(self: Self, other: SparkLikeExpr) -> Self:
+    def __rpow__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: other.__pow__(_input), "__rpow__", other=other
         ).alias("literal")
 
-    def __mod__(self: Self, other: SparkLikeExpr) -> Self:
+    def __mod__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__mod__(other), "__mod__", other=other
         )
 
-    def __rmod__(self: Self, other: SparkLikeExpr) -> Self:
+    def __rmod__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: other.__mod__(_input), "__rmod__", other=other
         ).alias("literal")
 
-    def __ge__(self: Self, other: SparkLikeExpr) -> Self:
+    def __ge__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__ge__(other), "__ge__", other=other
         )
 
-    def __gt__(self: Self, other: SparkLikeExpr) -> Self:
+    def __gt__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input > other, "__gt__", other=other
         )
 
-    def __le__(self: Self, other: SparkLikeExpr) -> Self:
+    def __le__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__le__(other), "__le__", other=other
         )
 
-    def __lt__(self: Self, other: SparkLikeExpr) -> Self:
+    def __lt__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__lt__(other), "__lt__", other=other
         )
 
-    def __and__(self: Self, other: SparkLikeExpr) -> Self:
+    def __and__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__and__(other), "__and__", other=other
         )
 
-    def __or__(self: Self, other: SparkLikeExpr) -> Self:
+    def __or__(self: Self, other: SparkLikeExpr[FrameT]) -> Self:
         return self._from_call(
             lambda _input, other: _input.__or__(other), "__or__", other=other
         )
@@ -524,14 +525,14 @@ def over(
         if (window_function := self._window_function) is not None:
             assert order_by is not None  # noqa: S101
 
-            def func(df: SparkLikeLazyFrame) -> list[Column]:
+            def func(df: SparkLikeLazyFrame[FrameT]) -> list[Column]:
                 return [
                     window_function(expr, partition_by, order_by)
                     for expr in self._call(df)
                 ]
         else:
 
-            def func(df: SparkLikeLazyFrame) -> list[Column]:
+            def func(df: SparkLikeLazyFrame[FrameT]) -> list[Column]:
                 return [
                     expr.over(self._Window.partitionBy(*partition_by))
                     for expr in self._call(df)

narwhals/_spark_like/group_by.py

@@ -1,18 +1,25 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from typing import Generic
+from typing import TypeVar
+from typing import cast
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from narwhals._spark_like.dataframe import DataFrame
     from narwhals._spark_like.dataframe import SparkLikeLazyFrame
+    from narwhals._spark_like.dataframe import SQLFrameDataFrame
     from narwhals._spark_like.expr import SparkLikeExpr
 
+FrameT = TypeVar("FrameT", "DataFrame", "SQLFrameDataFrame")
 
-class SparkLikeLazyGroupBy:
+
+class SparkLikeLazyGroupBy(Generic[FrameT]):
     def __init__(
         self: Self,
-        compliant_frame: SparkLikeLazyFrame,
+        compliant_frame: SparkLikeLazyFrame[FrameT],
         keys: list[str],
         drop_null_keys: bool,  # noqa: FBT001
     ) -> None:
@@ -22,7 +29,7 @@ def __init__(
             self._compliant_frame = compliant_frame
         self._keys = keys
 
-    def agg(self: Self, *exprs: SparkLikeExpr) -> SparkLikeLazyFrame:
+    def agg(self: Self, *exprs: SparkLikeExpr[FrameT]) -> SparkLikeLazyFrame[FrameT]:
         agg_columns = []
         df = self._compliant_frame
         for expr in exprs:
@@ -52,6 +59,7 @@ def agg(self: Self, *exprs: SparkLikeExpr) -> SparkLikeLazyFrame:
             return self._compliant_frame._from_native_frame(
                 self._compliant_frame._native_frame.select(*self._keys).dropDuplicates()
             )
-        return self._compliant_frame._from_native_frame(
-            self._compliant_frame._native_frame.groupBy(*self._keys).agg(*agg_columns)
+        native = self._compliant_frame._native_frame.groupBy(*self._keys).agg(
+            *agg_columns
         )
+        return self._compliant_frame._from_native_frame(cast("FrameT", native))