feat(expr-ir): The big Selector overhaul

narwhals/_plan/__init__.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+from narwhals._plan import selectors
 from narwhals._plan.dataframe import DataFrame
-from narwhals._plan.expr import Expr, Selector
-from narwhals._plan.expressions import selectors
+from narwhals._plan.expr import Expr
 from narwhals._plan.functions import (
     all,
     all_horizontal,
@@ -25,6 +25,7 @@
     sum_horizontal,
     when,
 )
+from narwhals._plan.selectors import Selector
 from narwhals._plan.series import Series
 
 __all__ = [

narwhals/_plan/_expr_ir.py

@@ -18,9 +18,10 @@
     from typing_extensions import Self
 
     from narwhals._plan.compliant.typing import Ctx, FrameT_contra, R_co
-    from narwhals._plan.expr import Expr, Selector
+    from narwhals._plan.expr import Expr
     from narwhals._plan.expressions.expr import Alias, Cast, Column
     from narwhals._plan.meta import MetaNamespace
+    from narwhals._plan.selectors import Selector
     from narwhals._plan.typing import ExprIRT2, MapIR, Seq
     from narwhals.dtypes import DType
 
@@ -191,11 +192,11 @@ def _map_ir_child(obj: ExprIR | Seq[ExprIR], fn: MapIR, /) -> ExprIR | Seq[ExprI
 
 class SelectorIR(ExprIR, config=ExprIROptions.no_dispatch()):
     def to_narwhals(self, version: Version = Version.MAIN) -> Selector:
-        from narwhals._plan import expr
+        from narwhals._plan.selectors import Selector, SelectorV1
 
         if version is Version.MAIN:
-            return expr.Selector._from_ir(self)
-        return expr.SelectorV1._from_ir(self)
+            return Selector._from_ir(self)
+        return SelectorV1._from_ir(self)
 
     def matches_column(self, name: str, dtype: DType) -> bool:
         """Return True if we can select this column.

narwhals/_plan/_guards.py

@@ -11,9 +11,10 @@
 if TYPE_CHECKING:
     from typing_extensions import TypeIs
 
-    from narwhals._plan import expr, expressions as ir
+    from narwhals._plan import expressions as ir
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.expr import Expr
+    from narwhals._plan.selectors import Selector
     from narwhals._plan.series import Series
     from narwhals._plan.typing import (
         ColumnNameOrSelector,
@@ -49,6 +50,12 @@ def _expr(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
     return expr
 
 
+def _selectors(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
+    from narwhals._plan import selectors
+
+    return selectors
+
+
 def _series(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
     from narwhals._plan import series
 
@@ -63,8 +70,8 @@ def is_expr(obj: Any) -> TypeIs[Expr]:
     return isinstance(obj, _expr().Expr)
 
 
-def is_selector(obj: Any) -> TypeIs[expr.Selector]:
-    return isinstance(obj, _expr().Selector)
+def is_selector(obj: Any) -> TypeIs[Selector]:
+    return isinstance(obj, _selectors().Selector)
 
 
 def is_column(obj: Any) -> TypeIs[Expr]:
@@ -83,7 +90,7 @@ def is_into_expr_column(obj: Any) -> TypeIs[IntoExprColumn]:
 def is_column_name_or_selector(
     obj: Any, *, allow_expr: bool = False
 ) -> TypeIs[ColumnNameOrSelector]:
-    tps = (str, _expr().Selector) if not allow_expr else (str, _expr().Expr)
+    tps = (str, _selectors().Selector) if not allow_expr else (str, _expr().Expr)
     return isinstance(obj, tps)
 
 

narwhals/_plan/expr.py

@@ -2,10 +2,10 @@
 
 import math
 from collections.abc import Iterable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, ClassVar, overload
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from narwhals._plan import common, expressions as ir
-from narwhals._plan._guards import is_column, is_expr, is_series
+from narwhals._plan._guards import is_expr, is_series
 from narwhals._plan._parse import (
     parse_into_expr_ir,
     parse_into_seq_of_expr_ir,
@@ -17,7 +17,6 @@
     functions as F,
     operators as ops,
 )
-from narwhals._plan.expressions.selectors import by_name
 from narwhals._plan.options import (
     EWMOptions,
     RankOptions,
@@ -29,7 +28,7 @@
 from narwhals.exceptions import ComputeError
 
 if TYPE_CHECKING:
-    from typing_extensions import Never, Self
+    from typing_extensions import Self
 
     from narwhals._plan._function import Function
     from narwhals._plan.expressions.categorical import ExprCatNamespace
@@ -51,8 +50,6 @@
     )
 
 
-# NOTE: Overly simplified placeholders for mocking typing
-# Entirely ignoring namespace + function binding
 class Expr:
     _ir: ir.ExprIR
     _version: ClassVar[Version] = Version.MAIN
@@ -83,9 +80,15 @@ def alias(self, name: str) -> Self:
     def cast(self, dtype: IntoDType) -> Self:
         return self._from_ir(self._ir.cast(common.into_dtype(dtype)))
 
+    # TODO @dangotbanned: Swap out with `exclude_s`
     def exclude(self, *names: OneOrIterable[str]) -> Self:
         return self._from_ir(ir.Exclude.from_names(self._ir, *names))
 
+    def exclude_s(self, *names: OneOrIterable[str]) -> Expr:
+        from narwhals._plan import selectors as cs
+
+        return (self.meta.as_selector() - cs.by_name(*names)).as_expr()
+
     def count(self) -> Self:
         return self._from_ir(agg.Count(expr=self._ir))
 
@@ -573,111 +576,5 @@ def str(self) -> ExprStringNamespace:
         return ExprStringNamespace(_expr=self)
 
 
-class Selector(Expr):
-    _ir: ir.SelectorIR
-
-    def __repr__(self) -> str:
-        return f"nw._plan.Selector({self.version.name.lower()}):\n{self._ir!r}"
-
-    @classmethod
-    def _from_ir(cls, selector_ir: ir.SelectorIR, /) -> Self:  # type: ignore[override]
-        obj = cls.__new__(cls)
-        obj._ir = selector_ir
-        return obj
-
-    def _to_expr(self) -> Expr:
-        return self._ir.to_narwhals(self.version)
-
-    @overload  # type: ignore[override]
-    def __or__(self, other: Self) -> Self: ...
-    @overload
-    def __or__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __or__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if isinstance(other, type(self)):
-            op = ops.Or()
-            return self._from_ir(op.to_binary_selector(self._ir, other._ir))
-        return self._to_expr() | other
-
-    @overload  # type: ignore[override]
-    def __and__(self, other: Self) -> Self: ...
-    @overload
-    def __and__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __and__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if is_column(other) and (name := other.meta.output_name()):
-            other = by_name(name)
-        if isinstance(other, type(self)):
-            op = ops.And()
-            return self._from_ir(op.to_binary_selector(self._ir, other._ir))
-        return self._to_expr() & other
-
-    @overload  # type: ignore[override]
-    def __sub__(self, other: Self) -> Self: ...
-    @overload
-    def __sub__(self, other: IntoExpr) -> Expr: ...
-    def __sub__(self, other: IntoExpr) -> Self | Expr:
-        if isinstance(other, type(self)):
-            op = ops.Sub()
-            return self._from_ir(op.to_binary_selector(self._ir, other._ir))
-        return self._to_expr() - other
-
-    @overload  # type: ignore[override]
-    def __xor__(self, other: Self) -> Self: ...
-    @overload
-    def __xor__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __xor__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if isinstance(other, type(self)):
-            op = ops.ExclusiveOr()
-            return self._from_ir(op.to_binary_selector(self._ir, other._ir))
-        return self._to_expr() ^ other
-
-    def __invert__(self) -> Self:
-        return self._from_ir(ir.InvertSelector(selector=self._ir))
-
-    def __add__(self, other: Any) -> Expr:  # type: ignore[override]
-        if isinstance(other, type(self)):
-            msg = "unsupported operand type(s) for op: ('Selector' + 'Selector')"
-            raise TypeError(msg)
-        return self._to_expr() + other  # type: ignore[no-any-return]
-
-    def __radd__(self, other: Any) -> Never:
-        msg = "unsupported operand type(s) for op: ('Expr' + 'Selector')"
-        raise TypeError(msg)
-
-    def __rsub__(self, other: Any) -> Never:
-        msg = "unsupported operand type(s) for op: ('Expr' - 'Selector')"
-        raise TypeError(msg)
-
-    @overload  # type: ignore[override]
-    def __rand__(self, other: Self) -> Self: ...
-    @overload
-    def __rand__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __rand__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if is_column(other) and (name := other.meta.output_name()):
-            return by_name(name) & self
-        return self._to_expr().__rand__(other)
-
-    @overload  # type: ignore[override]
-    def __ror__(self, other: Self) -> Self: ...
-    @overload
-    def __ror__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __ror__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if is_column(other) and (name := other.meta.output_name()):
-            return by_name(name) | self
-        return self._to_expr().__ror__(other)
-
-    @overload  # type: ignore[override]
-    def __rxor__(self, other: Self) -> Self: ...
-    @overload
-    def __rxor__(self, other: IntoExprColumn | int | bool) -> Expr: ...
-    def __rxor__(self, other: IntoExprColumn | int | bool) -> Self | Expr:
-        if is_column(other) and (name := other.meta.output_name()):
-            return by_name(name) ^ self
-        return self._to_expr().__rxor__(other)
-
-
 class ExprV1(Expr):
     _version: ClassVar[Version] = Version.V1
-
-
-class SelectorV1(Selector):
-    _version: ClassVar[Version] = Version.V1

narwhals/_plan/expressions/expr.py

@@ -68,14 +68,17 @@ def col(name: str, /) -> Column:
     return Column(name=name)
 
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 def cols(*names: str) -> Columns:
     return Columns(names=names)
 
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 def nth(index: int, /) -> Nth:
     return Nth(index=index)
 
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 def index_columns(*indices: int) -> IndexColumns:
     return IndexColumns(indices=indices)
 
@@ -104,10 +107,12 @@ def to_selector_ir(self) -> RootSelector:
         return cs.ByName.from_name(self.name).to_selector_ir()
 
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class _ColumnSelection(ExprIR, config=ExprIROptions.no_dispatch()):
     """Nodes which can resolve to `Column`(s) with a `Schema`."""
 
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class Columns(_ColumnSelection):
     __slots__ = ("names",)
     names: Seq[str]
@@ -119,24 +124,31 @@ def to_selector_ir(self) -> RootSelector:
         return cs.ByName.from_names(*self.names).to_selector_ir()
 
 
-# TODO @dangotbanned: Add `selectors.by_index`
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class Nth(_ColumnSelection):
     __slots__ = ("index",)
     index: int
 
     def __repr__(self) -> str:
         return f"nth({self.index})"
 
+    def to_selector_ir(self) -> RootSelector:
+        return cs.ByIndex.from_index(self.index).to_selector_ir()
+
 
-# TODO @dangotbanned: Add `selectors.by_index`
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class IndexColumns(_ColumnSelection):
     __slots__ = ("indices",)
     indices: Seq[int]
 
     def __repr__(self) -> str:
         return f"index_columns({self.indices!r})"
 
+    def to_selector_ir(self) -> RootSelector:
+        return cs.ByIndex.from_indices(self.indices).to_selector_ir()
+
 
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class All(_ColumnSelection):
     def __repr__(self) -> str:
         return "all()"
@@ -145,7 +157,7 @@ def to_selector_ir(self) -> RootSelector:
         return cs.All().to_selector_ir()
 
 
-# TODO @dangotbanned: Add `selectors.exclude`
+# TODO @dangotbanned: Scheduled to be removed, not needed with new selectors
 class Exclude(_ColumnSelection, child=("expr",)):
     __slots__ = ("expr", "names")
     expr: ExprIR