feat(typing): Generic BinaryExpr, default TypeVars

- Needed to add the `ruff` config in the end (#2578)
- Prevented forward refs in `TypeVar` from showing as unused imports

Author: dangotbanned

narwhals/_plan/expr.py

@@ -13,24 +13,18 @@
 import typing as t
 
 from narwhals._plan.common import ExprIR
+from narwhals._plan.typing import FunctionT, LeftT, OperatorT, RightT, RollingT
 
 if t.TYPE_CHECKING:
     from typing_extensions import Self
 
-    from narwhals._plan.common import Function, Seq
-    from narwhals._plan.functions import (
-        MapBatches,  # noqa: F401
-        RollingWindow,
-    )
+    from narwhals._plan.common import Seq
+    from narwhals._plan.functions import MapBatches  # noqa: F401
     from narwhals._plan.literal import LiteralValue
-    from narwhals._plan.operators import Operator
     from narwhals._plan.options import FunctionOptions, SortMultipleOptions, SortOptions
     from narwhals._plan.window import Window
     from narwhals.dtypes import DType
 
-_FunctionT = t.TypeVar("_FunctionT", bound="Function")
-_RollingT = t.TypeVar("_RollingT", bound="RollingWindow")
-
 
 class Alias(ExprIR):
     __slots__ = ("expr", "name")
@@ -95,7 +89,7 @@ def __repr__(self) -> str:
         return f"lit({self.value!r})"
 
 
-class BinaryExpr(ExprIR):
+class BinaryExpr(ExprIR, t.Generic[LeftT, OperatorT, RightT]):
     """Application of two exprs via an `Operator`.
 
     This ✅
@@ -107,9 +101,9 @@ class BinaryExpr(ExprIR):
 
     __slots__ = ("left", "op", "right")
 
-    left: ExprIR
-    op: Operator
-    right: ExprIR
+    left: LeftT
+    op: OperatorT
+    right: RightT
 
     @property
     def is_scalar(self) -> bool:
@@ -203,7 +197,7 @@ def iter_right(self) -> t.Iterator[ExprIR]:
         yield from self.expr.iter_right()
 
 
-class FunctionExpr(ExprIR, t.Generic[_FunctionT]):
+class FunctionExpr(ExprIR, t.Generic[FunctionT]):
     """**Representing `Expr::Function`**.
 
     https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/expr.rs#L114-L120
@@ -214,7 +208,7 @@ class FunctionExpr(ExprIR, t.Generic[_FunctionT]):
     __slots__ = ("function", "input", "options")
 
     input: Seq[ExprIR]
-    function: _FunctionT
+    function: FunctionT
     """Enum type is named `FunctionExpr` in `polars`.
 
     Mirroring *exactly* doesn't make much sense in OOP.
@@ -257,7 +251,7 @@ def iter_right(self) -> t.Iterator[ExprIR]:
             yield from e.iter_right()
 
 
-class RollingExpr(FunctionExpr[_RollingT]): ...
+class RollingExpr(FunctionExpr[RollingT]): ...
 
 
 class AnonymousExpr(FunctionExpr["MapBatches"]):

narwhals/_plan/operators.py

@@ -3,9 +3,12 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from narwhals._plan.expr import BinaryExpr
+    from narwhals._plan.typing import LeftT, RightT
 
-from narwhals._plan.common import ExprIR, Immutable
+from narwhals._plan.common import Immutable
 
 
 class Operator(Immutable):
@@ -28,10 +31,13 @@ def __repr__(self) -> str:
             Modulus: "%",
             And: "&",
             Or: "|",
+            ExclusiveOr: "^",
         }
         return m[tp]
 
-    def to_binary_expr(self, left: ExprIR, right: ExprIR, /) -> BinaryExpr:
+    def to_binary_expr(
+        self, left: LeftT, right: RightT, /
+    ) -> BinaryExpr[LeftT, Self, RightT]:
         from narwhals._plan.expr import BinaryExpr
 
         return BinaryExpr(left=left, op=self, right=right)
@@ -77,3 +83,6 @@ class And(Operator): ...
 
 
 class Or(Operator): ...
+
+
+class ExclusiveOr(Operator): ...

narwhals/_plan/typing.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import typing as t
+
+from narwhals._typing_compat import TypeVar
+
+if t.TYPE_CHECKING:
+    from narwhals._plan import operators as ops
+    from narwhals._plan.common import ExprIR, Function
+    from narwhals._plan.functions import RollingWindow
+
+__all__ = ["FunctionT", "LeftT", "OperatorT", "RightT", "RollingT"]
+
+
+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")

pyproject.toml

@@ -130,6 +130,7 @@ extend-exclude = ["**/this.py"]
 [tool.ruff.lint]
 preview = true
 explicit-preview-rules = true
+typing-modules = ["narwhals._typing_compat"]
 
 extend-safe-fixes = [
   "C419",    # unnecessary-comprehension-in-call