feat(RFC): A richer Expr IR

.pre-commit-config.yaml

@@ -19,7 +19,7 @@ repos:
   hooks:
     - id: codespell
       files: \.(py|rst|md)$
-      args: [--ignore-words-list=ser]
+      args: [--ignore-words-list=ser, --ignore-words-list=RightT]
       exclude: ^docs/api-completeness.md$
 - repo: https://github.com/crate-ci/typos
   rev: 'v1.39.0'
@@ -83,6 +83,8 @@ repos:
       entry: "self: Self"
       language: pygrep
       files: ^narwhals/
+      # mypy needs `Self` for `ExprIR.dispatch`
+      exclude: ^narwhals/_plan/.*\.py
     - id: dtypes-import
       name: don't import from narwhals.dtypes (use `Version.dtypes` instead)
       entry: |
@@ -101,7 +103,10 @@ repos:
             narwhals/_utils\.py|
             narwhals/stable/v./_?dtypes.py|
             narwhals/.*__init__.py|
-            narwhals/.*typing\.py
+            narwhals/.*typing\.py|
+            narwhals/_plan/functions\.py|
+            narwhals/_plan/expressions/ranges\.py|
+            narwhals/_plan/schema\.py
           )
     - id: pull-request-target
       name: don't use `pull_request_target`
@@ -123,6 +128,7 @@ repos:
         (?x)
           ^(tests/utils\.py)
           |^(tests/test_plugin/)
+          |^(tests/plan/utils\.py)
   -   id: no-commit-to-branch
   -   id: end-of-file-fixer
       exclude: .svg$

narwhals/_plan/__init__.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from narwhals._plan.dataframe import DataFrame
+from narwhals._plan.expr import Expr, Selector
+from narwhals._plan.expressions import selectors
+from narwhals._plan.functions import (
+    all,
+    all_horizontal,
+    any_horizontal,
+    col,
+    concat_str,
+    exclude,
+    int_range,
+    len,
+    lit,
+    max,
+    max_horizontal,
+    mean,
+    mean_horizontal,
+    median,
+    min,
+    min_horizontal,
+    nth,
+    sum,
+    sum_horizontal,
+    when,
+)
+from narwhals._plan.series import Series
+
+__all__ = [
+    "DataFrame",
+    "Expr",
+    "Selector",
+    "Series",
+    "all",
+    "all_horizontal",
+    "any_horizontal",
+    "col",
+    "concat_str",
+    "exclude",
+    "int_range",
+    "len",
+    "lit",
+    "max",
+    "max_horizontal",
+    "mean",
+    "mean_horizontal",
+    "median",
+    "min",
+    "min_horizontal",
+    "nth",
+    "selectors",
+    "sum",
+    "sum_horizontal",
+    "when",
+]

narwhals/_plan/_dispatch.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from operator import attrgetter
+from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, final, overload
+
+from narwhals._plan._guards import is_function_expr
+from narwhals._plan.compliant.typing import FrameT_contra, R_co
+from narwhals._typing_compat import TypeVar
+
+if TYPE_CHECKING:
+    from typing_extensions import Never, TypeAlias
+
+    from narwhals._plan.compliant.typing import Ctx
+    from narwhals._plan.expressions import ExprIR, FunctionExpr
+    from narwhals._plan.typing import ExprIRT, FunctionT
+
+__all__ = ["Dispatcher", "get_dispatch_name"]
+
+
+Node = TypeVar("Node", bound="ExprIR | FunctionExpr[Any]")
+Node_contra = TypeVar(
+    "Node_contra", bound="ExprIR | FunctionExpr[Any]", contravariant=True
+)
+Raiser: TypeAlias = Callable[..., "Never"]
+
+
+class Binder(Protocol[Node_contra]):
+    def __call__(
+        self, ctx: Ctx[FrameT_contra, R_co], /
+    ) -> BoundMethod[Node_contra, FrameT_contra, R_co]: ...
+
+
+class BoundMethod(Protocol[Node_contra, FrameT_contra, R_co]):
+    def __call__(self, node: Node_contra, frame: FrameT_contra, name: str, /) -> R_co: ...
+
+
+@final
+class Dispatcher(Generic[Node]):
+    """Translate class definitions into error-wrapped method calls.
+
+    Operates over `ExprIR` and `Function` nodes.
+
+    By default, we dispatch to the compliant-level by calling a method that is the
+    **snake_case**-equivalent of the class name:
+
+        class BinaryExpr(ExprIR): ...
+
+        class CompliantExpr(Protocol):
+            def binary_expr(self, *args: Any): ...
+    """
+
+    __slots__ = ("_bind", "_name")
+    _bind: Binder[Node]
+    _name: str
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}<{self.name}>"
+
+    def bind(
+        self, ctx: Ctx[FrameT_contra, R_co], /
+    ) -> BoundMethod[Node, FrameT_contra, R_co]:
+        """Retrieve the implementation of this expression from `ctx`.
+
+        Binds an instance method, most commonly via:
+
+            expr: CompliantExpr
+            method = getattr(expr, "method_name")
+        """
+        try:
+            return self._bind(ctx)
+        except AttributeError:
+            raise self._not_implemented_error(ctx, "compliant") from None
+
+    def __call__(
+        self,
+        node: Node,
+        ctx: Ctx[FrameT_contra, R_co],
+        frame: FrameT_contra,
+        name: str,
+        /,
+    ) -> R_co:
+        """Evaluate this expression in `frame`, using implementation(s) provided by `ctx`."""
+        method = self.bind(ctx)
+        if result := method(node, frame, name):
+            return result
+        raise self._not_implemented_error(ctx, "context")
+
+    @staticmethod
+    def from_expr_ir(tp: type[ExprIRT], /) -> Dispatcher[ExprIRT]:
+        if not tp.__expr_ir_config__.allow_dispatch:
+            return Dispatcher._no_dispatch(tp)
+        return Dispatcher._from_type(tp)
+
+    @staticmethod
+    def from_function(tp: type[FunctionT], /) -> Dispatcher[FunctionExpr[FunctionT]]:
+        return Dispatcher._from_type(tp)
+
+    @staticmethod
+    @overload
+    def _from_type(tp: type[ExprIRT], /) -> Dispatcher[ExprIRT]: ...
+    @staticmethod
+    @overload
+    def _from_type(tp: type[FunctionT], /) -> Dispatcher[FunctionExpr[FunctionT]]: ...
+    @staticmethod
+    def _from_type(tp: type[ExprIRT | FunctionT], /) -> Dispatcher[Any]:
+        obj = Dispatcher.__new__(Dispatcher)
+        obj._name = _method_name(tp)
+        getter = attrgetter(obj._name)
+        is_namespaced = tp.__expr_ir_config__.is_namespaced
+        obj._bind = _via_namespace(getter) if is_namespaced else getter
+        return obj
+
+    @staticmethod
+    def _no_dispatch(tp: type[ExprIRT], /) -> Dispatcher[ExprIRT]:
+        obj = Dispatcher.__new__(Dispatcher)
+        obj._name = tp.__name__
+        obj._bind = obj._make_no_dispatch_error()
+        return obj
+
+    def _make_no_dispatch_error(self) -> Callable[[Any], Raiser]:
+        def _no_dispatch_error(node: Node, *_: Any) -> Never:
+            msg = (
+                f"{self.name!r} should not appear at the compliant-level.\n\n"
+                f"Make sure to expand all expressions first, got:\n{node!r}"
+            )
+            raise TypeError(msg)
+
+        def getter(_: Any, /) -> Raiser:
+            return _no_dispatch_error
+
+        return getter
+
+    def _not_implemented_error(
+        self, ctx: object, /, missing: Literal["compliant", "context"]
+    ) -> NotImplementedError:
+        if missing == "context":
+            msg = f"`{self.name}` is not yet implemented for {type(ctx).__name__!r}"
+        else:
+            msg = (
+                f"`{self.name}` has not been implemented at the compliant-level.\n"
+                f"Hint: Try adding `CompliantExpr.{self.name}()` or `CompliantNamespace.{self.name}()`"
+            )
+        return NotImplementedError(msg)
+
+
+def _via_namespace(getter: Callable[[Any], Any], /) -> Callable[[Any], Any]:
+    def _(ctx: Any, /) -> Any:
+        return getter(ctx.__narwhals_namespace__())
+
+    return _
+
+
+def _pascal_to_snake_case(s: str) -> str:
+    """Convert a PascalCase string to snake_case.
+
+    Adapted from https://github.com/pydantic/pydantic/blob/f7a9b73517afecf25bf898e3b5f591dffe669778/pydantic/alias_generators.py#L43-L62
+    """
+    # Handle the sequence of uppercase letters followed by a lowercase letter
+    snake = _PATTERN_UPPER_LOWER.sub(_re_repl_snake, s)
+    # Insert an underscore between a lowercase letter and an uppercase letter
+    return _PATTERN_LOWER_UPPER.sub(_re_repl_snake, snake).lower()
+
+
+_PATTERN_UPPER_LOWER = re.compile(r"([A-Z]+)([A-Z][a-z])")
+_PATTERN_LOWER_UPPER = re.compile(r"([a-z])([A-Z])")
+
+
+def _re_repl_snake(match: re.Match[str], /) -> str:
+    return f"{match.group(1)}_{match.group(2)}"
+
+
+def _method_name(tp: type[ExprIRT | FunctionT]) -> str:
+    config = tp.__expr_ir_config__
+    name = config.override_name or _pascal_to_snake_case(tp.__name__)
+    return f"{ns}.{name}" if (ns := getattr(config, "accessor_name", "")) else name
+
+
+def get_dispatch_name(expr: ExprIR, /) -> str:
+    """Return the synthesized method name for `expr`."""
+    return (
+        repr(expr.function) if is_function_expr(expr) else expr.__expr_ir_dispatch__.name
+    )