chore: More coverage, simplify

Author: dangotbanned

narwhals/_plan/_expansion.py

@@ -191,7 +191,7 @@ def prepare_projection(self, exprs: Collection[ExprIR], /) -> Seq[NamedIR]:
                 raise NotImplementedError(msg)
             output_names.append(name)
             named_irs.append(ir.named_ir(name, remove_alias(target)))
-            root_names.update(meta._expr_to_leaf_column_names_iter(e))
+            root_names.update(meta.iter_root_names(e))
         if len(output_names) != len(set(output_names)):
             raise duplicate_error(expanded)
         if not (set(self.schema).issuperset(root_names)):

narwhals/_plan/meta.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 from functools import lru_cache
-from itertools import chain
 from typing import TYPE_CHECKING, Literal, overload
 
 from narwhals._plan import expressions as ir
@@ -19,7 +18,7 @@
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator
+    from collections.abc import Iterator
 
     from narwhals._plan import Selector
 
@@ -78,7 +77,7 @@ def output_name(self, *, raise_if_undetermined: bool = True) -> str | None:
 
     def root_names(self) -> list[str]:
         """Get the root column names."""
-        return list(_expr_to_leaf_column_names_iter(self._ir))
+        return list(iter_root_names(self._ir))
 
     def as_selector(self) -> Selector:
         """Try to turn this expression into a selector.
@@ -91,46 +90,17 @@ def as_selector(self) -> Selector:
         return self._ir.to_selector_ir().to_narwhals()
 
 
-def _expr_to_leaf_column_names_iter(expr: ir.ExprIR, /) -> Iterator[str]:
-    for e in _expr_to_leaf_column_exprs_iter(expr):
-        result = _expr_to_leaf_column_name(e)
-        if isinstance(result, str):
-            yield result
-
-
-def _expr_to_leaf_column_exprs_iter(expr: ir.ExprIR, /) -> Iterator[ir.ExprIR]:
-    for outer in expr.iter_root_names():
-        if isinstance(outer, ir.Column):
-            yield outer
-
-
-def _expr_to_leaf_column_name(expr: ir.ExprIR, /) -> str | ComputeError:
-    leaves = list(_expr_to_leaf_column_exprs_iter(expr))
-    if not len(leaves) <= 1:
-        msg = "found more than one root column name"
-        return ComputeError(msg)
-    if not leaves:
-        msg = "no root column name found"
-        return ComputeError(msg)
-    leaf = leaves[0]
-    if isinstance(leaf, ir.Column):
-        return leaf.name
-    msg = f"Expected unreachable, got {type(leaf).__name__!r}\n\n{leaf}"
-    return ComputeError(msg)
+def iter_root_names(expr: ir.ExprIR, /) -> Iterator[str]:
+    yield from (e.name for e in expr.iter_root_names() if isinstance(e, ir.Column))
 
 
 def root_name_first(expr: ir.ExprIR, /) -> str:
-    it = _expr_to_leaf_column_names_iter(expr)
-    if name := next(it, None):
+    if name := next(iter_root_names(expr), None):
         return name
-    msg = "`name.keep_name` expected at least one column name"
+    msg = f"`name.keep_name` expected at least one column name, got `{expr!r}`"
     raise InvalidOperationError(msg)
 
 
-def root_names_unique(exprs: Iterable[ir.ExprIR], /) -> set[str]:
-    return set(chain.from_iterable(_expr_to_leaf_column_names_iter(e) for e in exprs))
-
-
 @lru_cache(maxsize=32)
 def _expr_output_name(expr: ir.ExprIR, /) -> str | ComputeError:
     for e in expr.iter_output_name():

tests/plan/expr_expansion_test.py

@@ -12,6 +12,7 @@
 from narwhals.exceptions import (
     ColumnNotFoundError,
     DuplicateError,
+    InvalidOperationError,
     MultiOutputExpressionError,
 )
 from tests.plan.utils import Frame, assert_expr_ir_equal, named_ir, re_compile
@@ -128,6 +129,14 @@ def test_special_aliases_single(expr: nwp.Expr, expected: str) -> None:
     df.assert_selects(expr, expected)
 
 
+def test_keep_name_no_names(df_1: Frame) -> None:
+    with pytest.raises(
+        InvalidOperationError,
+        match=r"expected at least one.+name.+got.+lit.+1.+name.keep",
+    ):
+        df_1.project(nwp.lit(1).name.keep())
+
+
 def alias_replace_guarded(name: str) -> MapIR:
     """Guards against repeatedly creating the same alias."""
 

tests/plan/meta_test.py

@@ -2,6 +2,7 @@
 
 import datetime as dt
 import re
+import string
 from typing import Any
 
 import pytest
@@ -56,6 +57,11 @@
         ),
         (nwp.all().mean(), pl.all().mean(), []),
         (nwp.all().mean().sort_by("d"), pl.all().mean().sort_by("d"), ["d"]),
+        (
+            nwp.all_horizontal(*string.ascii_letters),
+            pl.all_horizontal(*string.ascii_letters),
+            list(string.ascii_letters),
+        ),
     ],
 )
 def test_meta_root_names(
@@ -319,6 +325,7 @@ def test_literal_output_name() -> None:
     assert e.meta.output_name() == ""
 
 
+# NOTE: Very low-priority
 @pytest.mark.xfail(
     reason="TODO: `Expr.struct.field` influences `meta.output_name`.",
     raises=AssertionError,