refactor: Move is_eager_expr and improve safety

- Now we don't accept any arbitrary expression
  - It must derive the current type
- Simplifies the typing
  - Really wasn't happy with what I had previously
  - Glad that is gone now
- Added a doc, since the single output part is really the most important

Author: dangotbanned

narwhals/_compliant/dataframe.py

@@ -27,9 +27,7 @@
     import pyarrow as pa
     from typing_extensions import Self
     from typing_extensions import TypeAlias
-    from typing_extensions import TypeIs
 
-    from narwhals._compliant.expr import EagerExpr
     from narwhals.dtypes import DType
     from narwhals.typing import SizeUnit
     from narwhals.typing import _2DArray
@@ -178,17 +176,16 @@ class EagerDataFrame(
     CompliantDataFrame[EagerSeriesT, EagerExprT_contra],
     Protocol[EagerSeriesT, EagerExprT_contra],
 ):
-    def _maybe_evaluate_expr(self, expr: EagerExprT_contra | T, /) -> EagerSeriesT | T:
-        if is_eager_expr(expr):
-            result: Sequence[EagerSeriesT] = expr(self)
-            if len(result) > 1:
-                msg = (
-                    "Multi-output expressions (e.g. `nw.all()` or `nw.col('a', 'b')`) "
-                    "are not supported in this context"
-                )
-                raise ValueError(msg)
-            return result[0]
-        return expr
+    def _evaluate_expr(self, expr: EagerExprT_contra, /) -> EagerSeriesT:
+        """Evaluate `expr` and ensure it has a **single** output."""
+        result: Sequence[EagerSeriesT] = expr(self)
+        if len(result) > 1:
+            msg = (
+                "Multi-output expressions (e.g. `nw.all()` or `nw.col('a', 'b')`) "
+                "are not supported in this context"
+            )
+            raise ValueError(msg)
+        return result[0]
 
     def _evaluate_into_exprs(self, *exprs: EagerExprT_contra) -> Sequence[EagerSeriesT]:
         return list(chain.from_iterable(self._evaluate_into_expr(expr) for expr in exprs))
@@ -209,11 +206,3 @@ def _evaluate_into_expr(self, expr: EagerExprT_contra, /) -> Sequence[EagerSerie
             msg = f"Safety assertion failed, expected {aliases}, got {result}"
             raise AssertionError(msg)
         return result
-
-
-# NOTE: `mypy` is requiring the gymnastics here and is very fragile
-# DON'T CHANGE THIS or `EagerDataFrame._maybe_evaluate_expr`
-def is_eager_expr(
-    obj: EagerExpr[Any, EagerSeriesT] | Any,
-) -> TypeIs[EagerExpr[Any, EagerSeriesT]]:
-    return hasattr(obj, "__narwhals_expr__")

narwhals/_compliant/expr.py

@@ -51,6 +51,7 @@
     from typing import Mapping
 
     from typing_extensions import Self
+    from typing_extensions import TypeIs
 
     from narwhals._compliant.namespace import CompliantNamespace
     from narwhals._compliant.namespace import EagerNamespace
@@ -390,6 +391,10 @@ def _reuse_series_extra_kwargs(
     ) -> dict[str, Any]:
         return {}
 
+    @classmethod
+    def _is_expr(cls, obj: Self | Any) -> TypeIs[Self]:
+        return isinstance(obj, cls)
+
     def _reuse_series_inner(
         self,
         df: EagerDataFrameT,
@@ -402,8 +407,8 @@ def _reuse_series_inner(
         kwargs = {
             **call_kwargs,
             **{
-                arg_name: df._maybe_evaluate_expr(arg_value)
-                for arg_name, arg_value in expressifiable_args.items()
+                name: df._evaluate_expr(value) if self._is_expr(value) else value
+                for name, value in expressifiable_args.items()
             },
         }
         method = methodcaller(