very wip selectors

None of this is functional yet

- Most of the upstream stuff is written in `python`
- The rust enum isn't very helpful for us
  - (`polars_plan::dsl::selector::Selector`)
  - It is opaque to the kind of selection

Author: dangotbanned

narwhals/_plan/dummy.py

@@ -27,7 +27,7 @@
 from narwhals.utils import Version, _hasattr_static, flatten
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
+    from typing_extensions import Never, Self
 
     from narwhals._plan.common import ExprIR, IntoExpr, IntoExprColumn, Seq, Udf
     from narwhals._plan.meta import ExprIRMetaNamespace
@@ -439,10 +439,77 @@ def meta(self) -> ExprIRMetaNamespace:
         return ExprIRMetaNamespace(self._ir)
 
 
+class DummySelector(DummyExpr):
+    _ir: expr.SelectorIR
+
+    @classmethod
+    def _from_ir(cls, ir: expr.SelectorIR, /) -> Self:  # type: ignore[override]
+        obj = cls.__new__(cls)
+        obj._ir = ir
+        return obj
+
+    def _to_expr(self) -> DummyExpr:
+        return self._ir.to_narwhals(self.version)
+
+    # TODO @dangotbanned: Make a decision on selector root, binary op
+    # Current typing warnings are accurate, this isn't valid yet
+    def __or__(self, other: t.Any) -> Self | t.Any:
+        if isinstance(other, type(self)):
+            op = ops.Or()
+            return self._from_ir(op.to_binary_selector(self._ir, other._ir))  # type: ignore[arg-type]
+        return self._to_expr() | other
+
+    def __and__(self, other: t.Any) -> Self | t.Any:
+        if isinstance(other, type(self)):
+            op = ops.And()
+            return self._from_ir(op.to_binary_selector(self._ir, other._ir))  # type: ignore[arg-type]
+        return self._to_expr() & other
+
+    def __sub__(self, other: t.Any) -> Self | t.Any:
+        if isinstance(other, type(self)):
+            op = ops.Sub()
+            return self._from_ir(op.to_binary_selector(self._ir, other._ir))  # type: ignore[arg-type]
+        return self._to_expr() - other
+
+    def __xor__(self, other: t.Any) -> Self | t.Any:
+        if isinstance(other, type(self)):
+            op = ops.ExclusiveOr()
+            return self._from_ir(op.to_binary_selector(self._ir, other._ir))  # type: ignore[arg-type]
+        return self._to_expr() ^ other
+
+    def __invert__(self) -> Never:
+        raise NotImplementedError
+
+    def __add__(self, other: t.Any) -> DummyExpr:  # 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 __rsub__(self, other: t.Any) -> Never:
+        raise NotImplementedError
+
+    def __rand__(self, other: t.Any) -> Never:
+        raise NotImplementedError
+
+    def __ror__(self, other: t.Any) -> Never:
+        raise NotImplementedError
+
+    def __rxor__(self, other: t.Any) -> Never:
+        raise NotImplementedError
+
+    def __radd__(self, other: t.Any) -> Never:
+        raise NotImplementedError
+
+
 class DummyExprV1(DummyExpr):
     _version: t.ClassVar[Version] = Version.V1
 
 
+class DummySelectorV1(DummySelector):
+    _version: t.ClassVar[Version] = Version.V1
+
+
 class DummyCompliantExpr:
     _ir: ExprIR
     _version: Version

narwhals/_plan/expr.py

@@ -13,15 +13,25 @@
 import typing as t
 
 from narwhals._plan.common import ExprIR
-from narwhals._plan.typing import FunctionT, LeftT, OperatorT, RightT, RollingT
+from narwhals._plan.typing import (
+    FunctionT,
+    LeftT,
+    OperatorT,
+    RightT,
+    RollingT,
+    SelectorOperatorT,
+)
+from narwhals.utils import Version
 
 if t.TYPE_CHECKING:
     from typing_extensions import Self
 
     from narwhals._plan.common import Seq
+    from narwhals._plan.dummy import DummySelector
     from narwhals._plan.functions import MapBatches  # noqa: F401
     from narwhals._plan.literal import LiteralValue
     from narwhals._plan.options import FunctionOptions, SortMultipleOptions, SortOptions
+    from narwhals._plan.selectors import Selector
     from narwhals._plan.window import Window
     from narwhals.dtypes import DType
 
@@ -408,9 +418,32 @@ def __repr__(self) -> str:
         return "*"
 
 
-class Selector(ExprIR):
+class SelectorIR(ExprIR):
+    """Not sure on this separation.
+
+    - Need a cleaner way of including `BinarySelector`.
+    - Like that there's easy access to operands
+    - Dislike that it inherits node iteration, since upstream doesn't use it for selectors
+    """
+
+    __slots__ = ("selector",)
+
+    selector: Selector
     """by_dtype, matches, numeric, boolean, string, categorical, datetime, all."""
 
+    def to_narwhals(self, version: Version = Version.MAIN) -> DummySelector:
+        from narwhals._plan import dummy
+
+        if version is Version.MAIN:
+            return dummy.DummySelector._from_ir(self)
+        return dummy.DummySelectorV1._from_ir(self)
+
+
+class BinarySelector(
+    BinaryExpr["SelectorIR", SelectorOperatorT, "SelectorIR"],
+    t.Generic[SelectorOperatorT],
+): ...
+
 
 class Ternary(ExprIR):
     """When-Then-Otherwise.

narwhals/_plan/meta.py

@@ -114,7 +114,7 @@ def _expr_output_name(ir: ExprIR) -> str | ComputeError:
 def _has_multiple_outputs(ir: ExprIR) -> bool:
     from narwhals._plan import expr
 
-    return isinstance(ir, (expr.Columns, expr.IndexColumns, expr.Selector, expr.All))
+    return isinstance(ir, (expr.Columns, expr.IndexColumns, expr.SelectorIR, expr.All))
 
 
 def _is_literal(ir: ExprIR, *, allow_aliasing: bool) -> bool:
@@ -145,7 +145,7 @@ def _is_column_selection(ir: ExprIR, *, allow_aliasing: bool) -> bool:
             expr.Exclude,
             expr.Nth,
             expr.IndexColumns,
-            expr.Selector,
+            expr.SelectorIR,
             expr.All,
         ),
     ):

narwhals/_plan/operators.py

@@ -5,7 +5,7 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from narwhals._plan.expr import BinaryExpr
+    from narwhals._plan.expr import BinaryExpr, BinarySelector, SelectorIR
     from narwhals._plan.typing import LeftT, RightT
 
 from narwhals._plan.common import Immutable
@@ -14,8 +14,8 @@
 class Operator(Immutable):
     def __repr__(self) -> str:
         tp = type(self)
-        if tp is Operator:
-            return "Operator"
+        if tp in {Operator, SelectorOperator}:
+            return tp.__name__
         m = {
             Eq: "==",
             NotEq: "!=",
@@ -43,6 +43,22 @@ def to_binary_expr(
         return BinaryExpr(left=left, op=self, right=right)
 
 
+class SelectorOperator(Operator):
+    """Operators that can *also* be used in selectors.
+
+    Remember that `Or` is named [`meta._selector_add`]!
+
+    [`meta._selector_add`]: https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/meta.rs#L113-L124
+    """
+
+    def to_binary_selector(
+        self, left: SelectorIR, right: SelectorIR, /
+    ) -> BinarySelector[Self]:
+        from narwhals._plan.expr import BinarySelector
+
+        return BinarySelector(left=left, op=self, right=right)
+
+
 class Eq(Operator): ...
 
 
@@ -64,7 +80,7 @@ class GtEq(Operator): ...
 class Add(Operator): ...
 
 
-class Sub(Operator): ...
+class Sub(SelectorOperator): ...
 
 
 class Multiply(Operator): ...
@@ -79,10 +95,10 @@ class FloorDivide(Operator): ...
 class Modulus(Operator): ...
 
 
-class And(Operator): ...
+class And(SelectorOperator): ...
 
 
-class Or(Operator): ...
+class Or(SelectorOperator): ...
 
 
-class ExclusiveOr(Operator): ...
+class ExclusiveOr(SelectorOperator): ...

narwhals/_plan/selectors.py

@@ -0,0 +1,144 @@
+"""Deviations from `polars`.
+
+- A `Selector` corresponds to a `nw.selectors` function
+- Binary ops are represented as a subtype of `BinaryExpr`
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Iterable
+
+from narwhals._plan.common import Immutable, is_iterable_reject
+from narwhals.utils import _parse_time_unit_and_time_zone
+
+if TYPE_CHECKING:
+    from datetime import timezone
+    from typing import Iterator, TypeVar
+
+    from narwhals._plan.dummy import DummySelector
+    from narwhals._plan.expr import SelectorIR
+    from narwhals.dtypes import DType
+    from narwhals.typing import TimeUnit
+
+    T = TypeVar("T")
+
+
+class Selector(Immutable):
+    def to_selector(self) -> SelectorIR:
+        from narwhals._plan.expr import SelectorIR
+
+        return SelectorIR(selector=self)
+
+
+class All(Selector): ...
+
+
+class ByDType(Selector):
+    __slots__ = ("dtypes",)
+
+    dtypes: frozenset[DType | type[DType]]
+
+    @staticmethod
+    def from_dtypes(
+        *dtypes: DType | type[DType] | Iterable[DType | type[DType]],
+    ) -> ByDType:
+        return ByDType(dtypes=frozenset(_flatten_hash_safe(dtypes)))
+
+
+class Boolean(Selector): ...
+
+
+class Categorical(Selector): ...
+
+
+class Datetime(Selector):
+    """Should swallow the [`utils` functions].
+
+    Just re-wrapping them for now, since `CompliantSelectorNamespace` is still using them.
+
+    [`utils` functions]: https://github.com/narwhals-dev/narwhals/blob/6d524ba04fca6fe2d6d25bdd69f75fabf1d79039/narwhals/utils.py#L1565-L1596
+    """
+
+    __slots__ = ("time_units", "time_zones")
+
+    time_units: frozenset[TimeUnit]
+    time_zones: frozenset[str | None]
+
+    @staticmethod
+    def from_time_unit_and_time_zone(
+        time_unit: TimeUnit | Iterable[TimeUnit] | None,
+        time_zone: str | timezone | Iterable[str | timezone | None] | None,
+        /,
+    ) -> Datetime:
+        units, zones = _parse_time_unit_and_time_zone(time_unit, time_zone)
+        return Datetime(time_units=frozenset(units), time_zones=frozenset(zones))
+
+
+class Matches(Selector):
+    __slots__ = ("pattern",)
+
+    pattern: re.Pattern[str]
+
+    @staticmethod
+    def from_string(pattern: str, /) -> Matches:
+        return Matches(pattern=re.compile(pattern))
+
+
+class Numeric(Selector): ...
+
+
+class String(Selector): ...
+
+
+def all() -> DummySelector:
+    return All().to_selector().to_narwhals()
+
+
+def by_dtype(
+    *dtypes: DType | type[DType] | Iterable[DType | type[DType]],
+) -> DummySelector:
+    return ByDType.from_dtypes(*dtypes).to_selector().to_narwhals()
+
+
+def boolean() -> DummySelector:
+    return Boolean().to_selector().to_narwhals()
+
+
+def categorical() -> DummySelector:
+    return Categorical().to_selector().to_narwhals()
+
+
+def datetime(
+    time_unit: TimeUnit | Iterable[TimeUnit] | None = None,
+    time_zone: str | timezone | Iterable[str | timezone | None] | None = ("*", None),
+) -> DummySelector:
+    return (
+        Datetime.from_time_unit_and_time_zone(time_unit, time_zone)
+        .to_selector()
+        .to_narwhals()
+    )
+
+
+def matches(pattern: str) -> DummySelector:
+    return Matches.from_string(pattern).to_selector().to_narwhals()
+
+
+def numeric() -> DummySelector:
+    return Numeric().to_selector().to_narwhals()
+
+
+def string() -> DummySelector:
+    return String().to_selector().to_narwhals()
+
+
+def _flatten_hash_safe(iterable: Iterable[T | Iterable[T]], /) -> Iterator[T]:
+    """Fully unwrap all levels of nesting.
+
+    Aiming to reduce the chances of passing an unhashable argument.
+    """
+    for element in iterable:
+        if isinstance(element, Iterable) and not is_iterable_reject(element):
+            yield from _flatten_hash_safe(element)
+        else:
+            yield element  # type: ignore[misc]

narwhals/_plan/typing.py

@@ -9,11 +9,15 @@
     from narwhals._plan.common import ExprIR, Function
     from narwhals._plan.functions import RollingWindow
 
-__all__ = ["FunctionT", "LeftT", "OperatorT", "RightT", "RollingT"]
+__all__ = ["FunctionT", "LeftT", "OperatorT", "RightT", "RollingT", "SelectorOperatorT"]
 
 
 FunctionT = TypeVar("FunctionT", bound="Function")
 RollingT = TypeVar("RollingT", bound="RollingWindow")
 LeftT = TypeVar("LeftT", bound="ExprIR", default="ExprIR")
 OperatorT = TypeVar("OperatorT", bound="ops.Operator", default="ops.Operator")
 RightT = TypeVar("RightT", bound="ExprIR", default="ExprIR")
+
+SelectorOperatorT = TypeVar(
+    "SelectorOperatorT", bound="ops.SelectorOperator", default="ops.SelectorOperator"
+)