feat: Support optional .otherwise(...)

Somewhat of typing nightmare but gets the job done for now

Author: dangotbanned

narwhals/_plan/demo.py

@@ -150,6 +150,10 @@ def when(*predicates: IntoExpr | t.Iterable[IntoExpr]) -> When:
         >>> when_then_many
         Narwhals DummyExpr (main):
         .when([(col('x')) == (lit(str: a))]).then(lit(int: 1)).otherwise(.when([(col('x')) == (lit(str: b))]).then(lit(int: 2)).otherwise(.when([(col('x')) == (lit(str: c))]).then(lit(int: 3)).otherwise(lit(int: 4))))
+        >>>
+        >>> nwd.when(nwd.col("y") == "b").then(1)
+        Narwhals DummyExpr (main):
+        .when([(col('y')) == (lit(str: b))]).then(lit(int: 1)).otherwise(lit(null))
     """
     if builtins.len(predicates) == 1 and is_expr(predicates[0]):
         expr = predicates[0]

narwhals/_plan/when_then.py

@@ -2,12 +2,12 @@
 
 from typing import TYPE_CHECKING
 
-from narwhals._plan.common import Immutable
+from narwhals._plan.common import Immutable, is_expr
+from narwhals._plan.dummy import DummyExpr
 from narwhals._plan.expr_parsing import parse_into_expr_ir
 
 if TYPE_CHECKING:
     from narwhals._plan.common import ExprIR, IntoExpr, Seq
-    from narwhals._plan.dummy import DummyExpr
     from narwhals._plan.expr import Ternary
 
 
@@ -24,7 +24,7 @@ def _from_expr(expr: DummyExpr, /) -> When:
         return When(condition=expr._ir)
 
 
-class Then(Immutable):
+class Then(Immutable, DummyExpr):
     __slots__ = ("condition", "statement")
 
     condition: ExprIR
@@ -37,9 +37,23 @@ def when(self, condition: IntoExpr, /) -> ChainedWhen:
         )
 
     def otherwise(self, statement: IntoExpr, /) -> DummyExpr:
-        return ternary_expr(
-            self.condition, self.condition, parse_into_expr_ir(statement)
-        ).to_narwhals()
+        return self._from_ir(self._otherwise(statement))
+
+    def _otherwise(self, statement: IntoExpr = None, /) -> ExprIR:
+        return ternary_expr(self.condition, self.statement, parse_into_expr_ir(statement))
+
+    @property
+    def _ir(self) -> ExprIR:  # type: ignore[override]
+        return self._otherwise()
+
+    @classmethod
+    def _from_ir(cls, ir: ExprIR, /) -> DummyExpr:  # type: ignore[override]
+        return DummyExpr._from_ir(ir)
+
+    def __eq__(self, value: object) -> DummyExpr | bool:  # type: ignore[override]
+        if is_expr(value):
+            return super(DummyExpr, self).__eq__(value)
+        return super().__eq__(value)
 
 
 class ChainedWhen(Immutable):
@@ -55,7 +69,7 @@ def then(self, statement: IntoExpr, /) -> ChainedThen:
         )
 
 
-class ChainedThen(Immutable):
+class ChainedThen(Immutable, DummyExpr):
     """https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/arity.rs#L89-L130."""
 
     __slots__ = ("conditions", "statements")
@@ -70,12 +84,28 @@ def when(self, condition: IntoExpr, /) -> ChainedWhen:
         )
 
     def otherwise(self, statement: IntoExpr, /) -> DummyExpr:
+        return self._from_ir(self._otherwise(statement))
+
+    def _otherwise(self, statement: IntoExpr = None, /) -> ExprIR:
         otherwise = parse_into_expr_ir(statement)
         it_conditions = reversed(self.conditions)
         it_statements = reversed(self.statements)
         for e in it_conditions:
             otherwise = ternary_expr(e, next(it_statements), otherwise)
-        return otherwise.to_narwhals()
+        return otherwise
+
+    @property
+    def _ir(self) -> ExprIR:  # type: ignore[override]
+        return self._otherwise()
+
+    @classmethod
+    def _from_ir(cls, ir: ExprIR, /) -> DummyExpr:  # type: ignore[override]
+        return DummyExpr._from_ir(ir)
+
+    def __eq__(self, value: object) -> DummyExpr | bool:  # type: ignore[override]
+        if is_expr(value):
+            return super(DummyExpr, self).__eq__(value)
+        return super().__eq__(value)
 
 
 def ternary_expr(predicate: ExprIR, truthy: ExprIR, falsy: ExprIR, /) -> Ternary: