feat: Spec'd-out CompliantExpr

narwhals/_arrow/expr.py

@@ -4,6 +4,7 @@
 from typing import Any
 from typing import Callable
 from typing import Literal
+from typing import Mapping
 from typing import Sequence
 
 from narwhals._arrow.expr_cat import ArrowExprCatNamespace
@@ -21,6 +22,7 @@
 from narwhals.exceptions import ColumnNotFoundError
 from narwhals.typing import CompliantExpr
 from narwhals.utils import Implementation
+from narwhals.utils import not_implemented
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -262,7 +264,7 @@ def var(self: Self, ddof: int) -> Self:
     def skew(self: Self) -> Self:
         return reuse_series_implementation(self, "skew", returns_scalar=True)
 
-    def cast(self: Self, dtype: DType) -> Self:
+    def cast(self: Self, dtype: DType | type[DType]) -> Self:
         return reuse_series_implementation(self, "cast", dtype=dtype)
 
     def abs(self: Self) -> Self:
@@ -385,7 +387,11 @@ def unique(self: Self) -> Self:
         return reuse_series_implementation(self, "unique", maintain_order=False)
 
     def replace_strict(
-        self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None
+        self: Self,
+        old: Sequence[Any] | Mapping[Any, Any],
+        new: Sequence[Any],
+        *,
+        return_dtype: DType | type[DType] | None,
     ) -> Self:
         return reuse_series_implementation(
             self, "replace_strict", old=old, new=new, return_dtype=return_dtype
@@ -417,7 +423,7 @@ def clip(self: Self, lower_bound: Any | None, upper_bound: Any | None) -> Self:
             self, "clip", lower_bound=lower_bound, upper_bound=upper_bound
         )
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
+    def over(self: Self, keys: Sequence[str], kind: ExprKind) -> Self:
         if not is_scalar_like(kind):
             msg = "Only aggregation or literal operations are supported in `over` context for PyArrow."
             raise NotImplementedError(msg)
@@ -434,8 +440,9 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
                 raise NotImplementedError(msg)
 
             tmp = df.group_by(*keys, drop_null_keys=False).agg(self)
+            on = list(keys)
             tmp = df.simple_select(*keys).join(
-                tmp, how="left", left_on=keys, right_on=keys, suffix="_right"
+                tmp, how="left", left_on=on, right_on=on, suffix="_right"
             )
             return [tmp[alias] for alias in aliases]
 
@@ -455,7 +462,7 @@ def mode(self: Self) -> Self:
     def map_batches(
         self: Self,
         function: Callable[[Any], Any],
-        return_dtype: DType | None,
+        return_dtype: DType | type[DType] | None,
     ) -> Self:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             input_series_list = self._call(df)
@@ -579,6 +586,8 @@ def rank(
             self, "rank", method=method, descending=descending
         )
 
+    ewm_mean = not_implemented()
+
     @property
     def dt(self: Self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)

narwhals/_arrow/series.py

@@ -519,7 +519,7 @@ def is_null(self: Self) -> Self:
     def is_nan(self: Self) -> Self:
         return self._from_native_series(pc.is_nan(self._native_series))
 
-    def cast(self: Self, dtype: DType) -> Self:
+    def cast(self: Self, dtype: DType | type[DType]) -> Self:
         ser = self._native_series
         data_type = narwhals_to_native_dtype(dtype, self._version)
         return self._from_native_series(pc.cast(ser, data_type))

narwhals/_dask/expr.py

@@ -23,6 +23,7 @@
 from narwhals.typing import CompliantExpr
 from narwhals.utils import Implementation
 from narwhals.utils import generate_temporary_column_name
+from narwhals.utils import not_implemented
 
 if TYPE_CHECKING:
     try:
@@ -392,12 +393,6 @@ def unique(self: Self) -> Self:
     def drop_nulls(self: Self) -> Self:
         return self._from_call(lambda _input: _input.dropna(), "drop_nulls")
 
-    def replace_strict(
-        self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None
-    ) -> Self:
-        msg = "`replace_strict` is not yet supported for Dask expressions"
-        raise NotImplementedError(msg)
-
     def abs(self: Self) -> Self:
         return self._from_call(lambda _input: _input.abs(), "abs")
 
@@ -538,7 +533,7 @@ def null_count(self: Self) -> Self:
             lambda _input: _input.isna().sum().to_series(), "null_count"
         )
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
+    def over(self: Self, keys: Sequence[str], kind: ExprKind) -> Self:
         # pandas is a required dependency of dask so it's safe to import this
         from narwhals._pandas_like.group_by import AGGREGATIONS_TO_PANDAS_EQUIVALENT
 
@@ -607,3 +602,24 @@ def dt(self: Self) -> DaskExprDateTimeNamespace:
     @property
     def name(self: Self) -> DaskExprNameNamespace:
         return DaskExprNameNamespace(self)
+
+    arg_min = not_implemented()
+    arg_max = not_implemented()
+    arg_true = not_implemented()
+    head = not_implemented()
+    tail = not_implemented()
+    mode = not_implemented()
+    sort = not_implemented()
+    rank = not_implemented()
+    sample = not_implemented()
+    map_batches = not_implemented()
+    ewm_mean = not_implemented()
+    rolling_sum = not_implemented()
+    rolling_mean = not_implemented()
+    rolling_var = not_implemented()
+    rolling_std = not_implemented()
+    gather_every = not_implemented()
+    replace_strict = not_implemented()
+
+    cat = not_implemented()  # pyright: ignore[reportAssignmentType]
+    list = not_implemented()  # pyright: ignore[reportAssignmentType]

narwhals/_duckdb/expr.py

@@ -24,6 +24,7 @@
 from narwhals._expression_parsing import ExprKind
 from narwhals.typing import CompliantExpr
 from narwhals.utils import Implementation
+from narwhals.utils import not_implemented
 
 if TYPE_CHECKING:
     import duckdb
@@ -482,3 +483,36 @@ def name(self: Self) -> DuckDBExprNameNamespace:
     @property
     def list(self: Self) -> DuckDBExprListNamespace:
         return DuckDBExprListNamespace(self)
+
+    arg_min = not_implemented()
+    arg_max = not_implemented()
+    arg_true = not_implemented()
+    head = not_implemented()
+    tail = not_implemented()
+    mode = not_implemented()
+    sort = not_implemented()
+    rank = not_implemented()
+    sample = not_implemented()
+    map_batches = not_implemented()
+    ewm_mean = not_implemented()
+    rolling_sum = not_implemented()
+    rolling_mean = not_implemented()
+    rolling_var = not_implemented()
+    rolling_std = not_implemented()
+    gather_every = not_implemented()
+    drop_nulls = not_implemented()
+    diff = not_implemented()
+    unique = not_implemented()
+    shift = not_implemented()
+    is_unique = not_implemented()
+    is_first_distinct = not_implemented()
+    is_last_distinct = not_implemented()
+    cum_sum = not_implemented()
+    cum_count = not_implemented()
+    cum_min = not_implemented()
+    cum_max = not_implemented()
+    cum_prod = not_implemented()
+    replace_strict = not_implemented()
+    over = not_implemented()
+
+    cat = not_implemented()  # pyright: ignore[reportAssignmentType]

narwhals/_pandas_like/expr.py

@@ -5,6 +5,7 @@
 from typing import Any
 from typing import Callable
 from typing import Literal
+from typing import Mapping
 from typing import Sequence
 
 from narwhals._expression_parsing import ExprKind
@@ -386,7 +387,11 @@ def drop_nulls(self: Self) -> Self:
         return reuse_series_implementation(self, "drop_nulls")
 
     def replace_strict(
-        self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None
+        self: Self,
+        old: Sequence[Any] | Mapping[Any, Any],
+        new: Sequence[Any],
+        *,
+        return_dtype: DType | type[DType] | None,
     ) -> Self:
         return reuse_series_implementation(
             self, "replace_strict", old=old, new=new, return_dtype=return_dtype
@@ -452,7 +457,7 @@ def alias_output_names(names: Sequence[str]) -> Sequence[str]:
             call_kwargs=self._call_kwargs,
         )
 
-    def over(self: Self, partition_by: list[str], kind: ExprKind) -> Self:
+    def over(self: Self, partition_by: Sequence[str], kind: ExprKind) -> Self:
         if not is_elementary_expression(self):
             msg = (
                 "Only elementary expressions are supported for `.over` in pandas-like backends.\n\n"
@@ -564,7 +569,7 @@ def mode(self: Self) -> Self:
     def map_batches(
         self: Self,
         function: Callable[[Any], Any],
-        return_dtype: DType | None,
+        return_dtype: DType | type[DType] | None,
     ) -> Self:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             input_series_list = self._call(df)

narwhals/_spark_like/expr.py

@@ -19,6 +19,7 @@
 from narwhals.dependencies import get_pyspark
 from narwhals.typing import CompliantExpr
 from narwhals.utils import Implementation
+from narwhals.utils import not_implemented
 from narwhals.utils import parse_version
 
 if TYPE_CHECKING:
@@ -486,7 +487,7 @@ def _n_unique(_input: Column) -> Column:
 
         return self._from_call(_n_unique, "n_unique")
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
+    def over(self: Self, keys: Sequence[str], kind: ExprKind) -> Self:
         def func(df: SparkLikeLazyFrame) -> list[Column]:
             return [expr.over(self._Window.partitionBy(*keys)) for expr in self._call(df)]
 
@@ -526,3 +527,36 @@ def dt(self: Self) -> SparkLikeExprDateTimeNamespace:
     @property
     def list(self: Self) -> SparkLikeExprListNamespace:
         return SparkLikeExprListNamespace(self)
+
+    arg_min = not_implemented()
+    arg_max = not_implemented()
+    arg_true = not_implemented()
+    head = not_implemented()
+    tail = not_implemented()
+    mode = not_implemented()
+    sort = not_implemented()
+    rank = not_implemented()
+    sample = not_implemented()
+    map_batches = not_implemented()
+    ewm_mean = not_implemented()
+    rolling_sum = not_implemented()
+    rolling_mean = not_implemented()
+    rolling_var = not_implemented()
+    rolling_std = not_implemented()
+    gather_every = not_implemented()
+    drop_nulls = not_implemented()
+    diff = not_implemented()
+    unique = not_implemented()
+    shift = not_implemented()
+    is_first_distinct = not_implemented()
+    is_last_distinct = not_implemented()
+    cum_sum = not_implemented()
+    cum_count = not_implemented()
+    cum_min = not_implemented()
+    cum_max = not_implemented()
+    cum_prod = not_implemented()
+    replace_strict = not_implemented()
+    fill_null = not_implemented()
+    quantile = not_implemented()
+
+    cat = not_implemented()  # pyright: ignore[reportAssignmentType]

narwhals/expr.py

@@ -31,13 +31,18 @@
     from typing_extensions import Concatenate
     from typing_extensions import ParamSpec
     from typing_extensions import Self
+    from typing_extensions import TypeAlias
 
     from narwhals.dtypes import DType
     from narwhals.typing import CompliantExpr
+    from narwhals.typing import CompliantNamespace
     from narwhals.typing import IntoExpr
 
     PS = ParamSpec("PS")
     R = TypeVar("R")
+    _ToCompliant: TypeAlias = Callable[
+        [CompliantNamespace[Any, Any]], CompliantExpr[Any, Any]
+    ]
 
 
 class Expr:
@@ -47,7 +52,7 @@ def __init__(
         metadata: ExprMetadata,
     ) -> None:
         # callable from CompliantNamespace to CompliantExpr
-        self._to_compliant_expr = to_compliant_expr
+        self._to_compliant_expr: _ToCompliant = to_compliant_expr
         self._metadata = metadata
 
     def _from_callable(self, to_compliant_expr: Callable[[Any], Any]) -> Self:
@@ -608,7 +613,7 @@ def var(self: Self, *, ddof: int = 1) -> Self:
 
     def map_batches(
         self: Self,
-        function: Callable[[Any], Self],
+        function: Callable[[Any], CompliantExpr[Any, Any]],
         return_dtype: DType | None = None,
     ) -> Self:
         """Apply a custom python function to a whole Series or sequence of Series.

narwhals/functions.py

@@ -1490,7 +1490,7 @@ def func(plx: CompliantNamespace[Any, Any]) -> CompliantExpr[Any, Any]:
             compliant_value = extract_compliant(plx, value, str_as_lit=False)
             if is_scalar_like(kind) and is_compliant_expr(compliant_value):
                 compliant_value = compliant_value.broadcast(kind)
-            return compliant_expr.otherwise(compliant_value)  # type: ignore[no-any-return]
+            return compliant_expr.otherwise(compliant_value)  # type: ignore[attr-defined, no-any-return]
 
         return Expr(
             func,