feat: Re-introduce error, display shape mismatch

Resolves #3233 (comment)

Author: dangotbanned

narwhals/_plan/_expansion.py

@@ -42,7 +42,11 @@
 from typing import TYPE_CHECKING, Any, Union
 
 from narwhals._plan import common, expressions as ir, meta
-from narwhals._plan.exceptions import column_not_found_error, duplicate_error
+from narwhals._plan.exceptions import (
+    binary_expr_multi_output_error,
+    column_not_found_error,
+    duplicate_error,
+)
 from narwhals._plan.expressions import (
     Alias,
     ExprIR,
@@ -315,8 +319,7 @@ def _expand_binary_expr(self, origin: ir.BinaryExpr, /) -> Iterator[ir.BinaryExp
             binary = origin.__replace__(right=rights[0])
             yield from (binary.__replace__(left=left) for left in lefts)
         else:
-            msg = "TODO: `binary_expr_multi_output_error`"
-            raise MultiOutputExpressionError(msg)
+            raise binary_expr_multi_output_error(origin, lefts, rights)
 
     def _expand_function_expr(
         self, origin: ir.FunctionExpr, /

narwhals/_plan/exceptions.py

@@ -54,35 +54,41 @@ def hist_bins_monotonic_error(bins: Seq[float]) -> ComputeError:  # noqa: ARG001
     return ComputeError(msg)
 
 
-# NOTE: Always underlining `right`, since the message refers to both types of exprs
-# Assuming the most recent as the issue
+def _binary_underline(
+    left: ir.ExprIR,
+    operator: Operator,
+    right: ir.ExprIR,
+    /,
+    *,
+    underline_right: bool = True,
+) -> str:
+    lhs, op, rhs = repr(left), repr(operator), repr(right)
+    if underline_right:
+        indent = (len(lhs) + len(op) + 2) * " "
+        underline = len(rhs) * "^"
+    else:
+        indent = ""
+        underline = len(lhs) * "^"
+    return f"{lhs} {op} {rhs}\n{indent}{underline}"
+
+
 def binary_expr_shape_error(
     left: ir.ExprIR, op: Operator, right: ir.ExprIR
 ) -> ShapeError:
-    lhs_op = f"{left!r} {op!r} "
-    rhs = repr(right)
-    indent = len(lhs_op) * " "
-    underline = len(rhs) * "^"
+    expr = _binary_underline(left, op, right, underline_right=True)
     msg = (
-        f"Cannot combine length-changing expressions with length-preserving ones.\n"
-        f"{lhs_op}{rhs}\n{indent}{underline}"
+        f"Cannot combine length-changing expressions with length-preserving ones.\n{expr}"
     )
     return ShapeError(msg)
 
 
-# TODO @dangotbanned: Share the right underline code w/ `binary_expr_shape_error`
 def binary_expr_multi_output_error(
-    left: ir.ExprIR, op: Operator, right: ir.ExprIR
+    origin: ir.BinaryExpr, left_expand: Seq[ir.ExprIR], right_expand: Seq[ir.ExprIR]
 ) -> MultiOutputExpressionError:
-    lhs_op = f"{left!r} {op!r} "
-    rhs = repr(right)
-    indent = len(lhs_op) * " "
-    underline = len(rhs) * "^"
-    msg = (
-        "Multi-output expressions are only supported on the "
-        f"left-hand side of a binary operation.\n"
-        f"{lhs_op}{rhs}\n{indent}{underline}"
-    )
+    len_left, len_right = len(left_expand), len(right_expand)
+    lhs, op, rhs = origin.left, origin.op, origin.right
+    expr = _binary_underline(lhs, op, rhs, underline_right=len_left < len_right)
+    msg = f"Cannot combine selectors that produce a different number of columns ({len_left} != {len_right}).\n{expr}"
     return MultiOutputExpressionError(msg)
 
 

tests/plan/expr_expansion_test.py

@@ -678,23 +678,42 @@ def test_expand_binary_expr_combination(
         assert_expr_ir_equal(actual, expect)
 
 
-@pytest.mark.xfail(reason="TODO: Move fancy error message", raises=AssertionError)
-def test_expand_binary_expr_combination_invalid(df_1: Frame) -> None:  # pragma: no cover
-    pattern = re.escape(
+def test_expand_binary_expr_combination_invalid(df_1: Frame) -> None:
+    # fmt: off
+    expr = re.escape(
         "ncs.all() + ncs.by_name('b', 'c', require_all=True)\n"
-        "            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
+        "^^^^^^^^^"
     )
+    # fmt: on
+    shapes = "(20 != 2)"
+    pattern = rf"{shapes}.+\n{expr}"
     all_to_two = nwp.all() + nwp.col("b", "c")
     with pytest.raises(MultiOutputExpressionError, match=pattern):
         df_1.project(all_to_two)
 
-    pattern = re.escape(
-        "ncs.by_name('a', 'b', 'c', require_all=True).abs().fill_null([lit(int: 0)]).round() * ncs.by_index([9, 10], require_all=True).cast(Int64).sort(asc)\n"
-        "                                                                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
+    expr = re.escape(
+        "ncs.by_name('a', 'b', require_all=True).abs().fill_null([lit(int: 0)]).round() * ncs.by_index([9, 10, 11], require_all=True).cast(Int64).sort(asc)\n"
+        "                                                                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
     )
-    three_to_two = (
-        nwp.col("a", "b", "c").abs().fill_null(0).round(2)
-        * nwp.nth(9, 10).cast(nw.Int64).sort()
+    shapes = "(2 != 3)"
+    pattern = rf"{shapes}.+\n{expr}"
+    two_to_three = (
+        nwp.col("a", "b").abs().fill_null(0).round(2)
+        * nwp.nth(9, 10, 11).cast(nw.Int64).sort()
     )
     with pytest.raises(MultiOutputExpressionError, match=pattern):
-        df_1.project(three_to_two)
+        df_1.project(two_to_three)
+
+    # fmt: off
+    expr = re.escape(
+        "ncs.numeric() / [(ncs.numeric()) - (ncs.by_dtype([Int64]))]\n"
+        "^^^^^^^^^^^^^"
+    )
+    # fmt: on
+    shapes = "(10 != 9)"
+    pattern = rf"{shapes}.+\n{expr}"
+    ten_to_nine = (
+        ncs.numeric().as_expr() / (ncs.numeric() - ncs.by_dtype(nw.Int64)).as_expr()
+    )
+    with pytest.raises(MultiOutputExpressionError, match=pattern):
+        df_1.project(ten_to_nine)