feat: Adds Compliant(When|Then)

narwhals/_arrow/namespace.py

@@ -5,7 +5,6 @@
 from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Callable
 from typing import Iterable
 from typing import Literal
 from typing import Sequence
@@ -24,23 +23,20 @@
 from narwhals._arrow.utils import nulls_like
 from narwhals._arrow.utils import vertical_concat
 from narwhals._compliant import EagerNamespace
+from narwhals._compliant.when_then import CompliantThen
+from narwhals._compliant.when_then import CompliantWhen
 from narwhals._expression_parsing import combine_alias_output_names
 from narwhals._expression_parsing import combine_evaluate_output_names
 from narwhals.utils import Implementation
 from narwhals.utils import import_dtypes_module
 
 if TYPE_CHECKING:
-    from typing import Callable
-
     from typing_extensions import Self
-    from typing_extensions import TypeAlias
 
     from narwhals._arrow.typing import Incomplete
     from narwhals.dtypes import DType
     from narwhals.utils import Version
 
-    _Scalar: TypeAlias = Any
-
 
 class ArrowNamespace(EagerNamespace[ArrowDataFrame, ArrowSeries, ArrowExpr]):
     @property
@@ -253,7 +249,7 @@ def selectors(self: Self) -> ArrowSelectorNamespace:
         return ArrowSelectorNamespace(self)
 
     def when(self: Self, predicate: ArrowExpr) -> ArrowWhen:
-        return ArrowWhen(predicate, self._backend_version, version=self._version)
+        return ArrowWhen(predicate, context=self)
 
     def concat_str(
         self: Self,
@@ -293,25 +289,14 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
         )
 
 
-class ArrowWhen:
-    def __init__(
-        self: Self,
-        condition: ArrowExpr,
-        backend_version: tuple[int, ...],
-        then_value: ArrowExpr | _Scalar = None,
-        otherwise_value: ArrowExpr | _Scalar = None,
-        *,
-        version: Version,
-    ) -> None:
-        self._backend_version = backend_version
-        self._condition: ArrowExpr = condition
-        self._then_value: ArrowExpr | _Scalar = then_value
-        self._otherwise_value: ArrowExpr | _Scalar = otherwise_value
-        self._version = version
+class ArrowWhen(CompliantWhen[ArrowDataFrame, ArrowSeries, ArrowExpr]):
+    @property
+    def _then(self) -> type[ArrowThen]:
+        return ArrowThen
 
-    def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
+    def __call__(self, df: ArrowDataFrame, /) -> Sequence[ArrowSeries]:
         condition = self._condition(df)[0]
-        condition_native = condition._native_series
+        condition_native = condition.native
 
         if isinstance(self._then_value, ArrowExpr):
             value_series = self._then_value(df)[0]
@@ -346,46 +331,5 @@ def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
             )
         ]
 
-    def then(self: Self, value: ArrowExpr | ArrowSeries | _Scalar) -> ArrowThen:
-        self._then_value = value
 
-        return ArrowThen(
-            self,
-            depth=0,
-            function_name="whenthen",
-            evaluate_output_names=getattr(
-                value, "_evaluate_output_names", lambda _df: ["literal"]
-            ),
-            alias_output_names=getattr(value, "_alias_output_names", None),
-            backend_version=self._backend_version,
-            version=self._version,
-        )
-
-
-class ArrowThen(ArrowExpr):
-    def __init__(
-        self: Self,
-        call: ArrowWhen,
-        *,
-        depth: int,
-        function_name: str,
-        evaluate_output_names: Callable[[ArrowDataFrame], Sequence[str]],
-        alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
-        backend_version: tuple[int, ...],
-        version: Version,
-        call_kwargs: dict[str, Any] | None = None,
-        implementation: Implementation | None = None,
-    ) -> None:
-        self._backend_version = backend_version
-        self._version = version
-        self._call: ArrowWhen = call
-        self._depth = depth
-        self._function_name = function_name
-        self._evaluate_output_names = evaluate_output_names
-        self._alias_output_names = alias_output_names
-        self._call_kwargs = call_kwargs or {}
-
-    def otherwise(self: Self, value: ArrowExpr | ArrowSeries | _Scalar) -> ArrowExpr:
-        self._call._otherwise_value = value
-        self._function_name = "whenotherwise"
-        return self
+class ArrowThen(CompliantThen[ArrowDataFrame, ArrowSeries, ArrowExpr], ArrowExpr): ...

narwhals/_compliant/typing.py

@@ -36,6 +36,9 @@
 NativeExprT_co = TypeVar("NativeExprT_co", bound="NativeExpr", covariant=True)
 NativeSeriesT_co = TypeVar("NativeSeriesT_co", bound="NativeSeries", covariant=True)
 CompliantSeriesT = TypeVar("CompliantSeriesT", bound="CompliantSeries[Any]")
+CompliantSeriesOrNativeExprT = TypeVar(
+    "CompliantSeriesOrNativeExprT", bound="CompliantSeries[Any] | NativeExpr"
+)
 CompliantSeriesOrNativeExprT_co = TypeVar(
     "CompliantSeriesOrNativeExprT_co",
     bound="CompliantSeries[Any] | NativeExpr",

narwhals/_compliant/when_then.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+from typing import Sequence
+from typing import cast
+
+from narwhals._compliant.expr import CompliantExpr
+from narwhals._compliant.typing import CompliantExprT
+from narwhals._compliant.typing import CompliantFrameT
+from narwhals._compliant.typing import CompliantSeriesOrNativeExprT
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+    from typing_extensions import TypeAlias
+
+    from narwhals.utils import Implementation
+    from narwhals.utils import Version
+    from narwhals.utils import _FullContext
+
+if not TYPE_CHECKING:  # pragma: no cover
+    if sys.version_info >= (3, 9):
+        from typing import Protocol as Protocol38
+    else:
+        from typing import Generic as Protocol38
+else:  # pragma: no cover
+    # TODO @dangotbanned: Remove after dropping `3.8` (#2084)
+    # - https://github.com/narwhals-dev/narwhals/pull/2064#discussion_r1965921386
+    from typing import Protocol as Protocol38
+
+_Scalar: TypeAlias = Any
+
+
+class CompliantWhen(
+    Protocol38[CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT]
+):
+    _condition: CompliantExprT
+    _then_value: CompliantExprT | CompliantSeriesOrNativeExprT | _Scalar
+    _otherwise_value: CompliantExprT | CompliantSeriesOrNativeExprT | _Scalar
+    _implementation: Implementation
+    _backend_version: tuple[int, ...]
+    _version: Version
+
+    @property
+    def _then(
+        self,
+    ) -> type[
+        CompliantThen[CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT]
+    ]: ...
+
+    def __call__(
+        self, compliant_frame: CompliantFrameT, /
+    ) -> Sequence[CompliantSeriesOrNativeExprT]: ...
+
+    def then(
+        self, value: CompliantExprT | CompliantSeriesOrNativeExprT | _Scalar, /
+    ) -> CompliantThen[CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT]:
+        return self._then.from_when(self, value)
+
+    def __init__(self, condition: CompliantExprT, /, *, context: _FullContext) -> None:
+        self._condition = condition
+        self._then_value = None
+        self._otherwise_value = None
+        self._implementation = context._implementation
+        self._backend_version = context._backend_version
+        self._version = context._version
+
+
+class CompliantThen(
+    CompliantExpr[CompliantFrameT, CompliantSeriesOrNativeExprT],
+    Protocol38[CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT],
+):
+    _call: Callable[[CompliantFrameT], Sequence[CompliantSeriesOrNativeExprT]]
+    _when: CompliantWhen[CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT]
+    _function_name: str
+    _implementation: Implementation
+    _backend_version: tuple[int, ...]
+    _version: Version
+    _call_kwargs: dict[str, Any]
+
+    @classmethod
+    def from_when(
+        cls,
+        when: CompliantWhen[
+            CompliantFrameT, CompliantSeriesOrNativeExprT, CompliantExprT
+        ],
+        then_value: CompliantExprT | CompliantSeriesOrNativeExprT | _Scalar,
+        /,
+    ) -> Self:
+        when._then_value = then_value
+        obj = cls.__new__(cls)
+        obj._call = when
+        obj._when = when
+        obj._depth = 0
+        obj._function_name = "whenthen"
+        obj._evaluate_output_names = getattr(
+            then_value, "_evaluate_output_names", lambda _df: ["literal"]
+        )
+        obj._alias_output_names = getattr(then_value, "_alias_output_names", None)
+        obj._implementation = when._implementation
+        obj._backend_version = when._backend_version
+        obj._version = when._version
+        obj._call_kwargs = {}
+        return obj
+
+    def otherwise(
+        self, value: CompliantExprT | CompliantSeriesOrNativeExprT | _Scalar, /
+    ) -> CompliantExprT:
+        self._when._otherwise_value = value
+        self._function_name = "whenotherwise"
+        return cast("CompliantExprT", self)