fix(expr-ir): Ensure only __slots__, and not __dict__ too (#3201)

Author: dangotbanned

narwhals/_plan/_immutable.py

@@ -1,44 +1,21 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, TypeVar
+from typing import TYPE_CHECKING
+
+# ruff: noqa: N806
+from narwhals._plan._meta import ImmutableMeta
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
-    from typing import Any, Callable
-
-    from typing_extensions import Never, Self, dataclass_transform
-
-else:
-    # https://docs.python.org/3/library/typing.html#typing.dataclass_transform
-    def dataclass_transform(
-        *,
-        eq_default: bool = True,
-        order_default: bool = False,
-        kw_only_default: bool = False,
-        frozen_default: bool = False,
-        field_specifiers: tuple[type[Any] | Callable[..., Any], ...] = (),
-        **kwargs: Any,
-    ) -> Callable[[T], T]:
-        def decorator(cls_or_fn: T) -> T:
-            cls_or_fn.__dataclass_transform__ = {
-                "eq_default": eq_default,
-                "order_default": order_default,
-                "kw_only_default": kw_only_default,
-                "frozen_default": frozen_default,
-                "field_specifiers": field_specifiers,
-                "kwargs": kwargs,
-            }
-            return cls_or_fn
-
-        return decorator
-
-
-T = TypeVar("T")
-_IMMUTABLE_HASH_NAME: Literal["__immutable_hash_value__"] = "__immutable_hash_value__"
-
-
-@dataclass_transform(kw_only_default=True, frozen_default=True)
-class Immutable:
+    from collections.abc import Iterable, Iterator
+    from typing import Any, ClassVar, Final
+
+    from typing_extensions import Never, Self
+
+
+_HASH_NAME: Final = "__immutable_hash_value__"
+
+
+class Immutable(metaclass=ImmutableMeta):
     """A poor man's frozen dataclass.
 
     - Keyword-only constructor (IDE supported)
@@ -49,40 +26,43 @@ class Immutable:
     [`copy.replace`]: https://docs.python.org/3.13/library/copy.html#copy.replace
     """
 
-    __slots__ = (_IMMUTABLE_HASH_NAME,)
-    __immutable_hash_value__: int
+    __slots__ = (_HASH_NAME,)
+    if not TYPE_CHECKING:
+        # NOTE: Trying to avoid this being added to synthesized `__init__`
+        # Seems to be the only difference when decorating the metaclass
+        __immutable_hash_value__: int
 
-    @property
-    def __immutable_keys__(self) -> Iterator[str]:
-        slots: tuple[str, ...] = self.__slots__
-        for name in slots:
-            if name != _IMMUTABLE_HASH_NAME:
-                yield name
+    __immutable_keys__: ClassVar[tuple[str, ...]]
 
     @property
     def __immutable_values__(self) -> Iterator[Any]:
+        """Override to configure hash seed."""
+        getattr_ = getattr
         for name in self.__immutable_keys__:
-            yield getattr(self, name)
+            yield getattr_(self, name)
 
     @property
     def __immutable_items__(self) -> Iterator[tuple[str, Any]]:
+        getattr_ = getattr
         for name in self.__immutable_keys__:
-            yield name, getattr(self, name)
+            yield name, getattr_(self, name)
 
     @property
     def __immutable_hash__(self) -> int:
-        if hasattr(self, _IMMUTABLE_HASH_NAME):
-            return self.__immutable_hash_value__
-        hash_value = hash((self.__class__, *self.__immutable_values__))
-        object.__setattr__(self, _IMMUTABLE_HASH_NAME, hash_value)
-        return self.__immutable_hash_value__
+        HASH = _HASH_NAME
+        if hasattr(self, HASH):
+            hash_value: int = getattr(self, HASH)
+        else:
+            hash_value = hash((self.__class__, *self.__immutable_values__))
+            object.__setattr__(self, HASH, hash_value)
+        return hash_value
 
     def __setattr__(self, name: str, value: Never) -> Never:
         msg = f"{type(self).__name__!r} is immutable, {name!r} cannot be set."
         raise AttributeError(msg)
 
     def __replace__(self, **changes: Any) -> Self:
-        """https://docs.python.org/3.13/library/copy.html#copy.replace"""  # noqa: D415
+        """https://docs.python.org/3.13/library/copy.html#copy.replace."""
         if len(changes) == 1:
             # The most common case is a single field replacement.
             # Iff that field happens to be equal, we can noop, preserving the current object's hash.
@@ -96,13 +76,6 @@ def __replace__(self, **changes: Any) -> Self:
                     changes[name] = value_current
         return type(self)(**changes)
 
-    def __init_subclass__(cls, *args: Any, **kwds: Any) -> None:
-        super().__init_subclass__(*args, **kwds)
-        if cls.__slots__:
-            ...
-        else:
-            cls.__slots__ = ()
-
     def __hash__(self) -> int:
         return self.__immutable_hash__
 
@@ -111,35 +84,26 @@ def __eq__(self, other: object) -> bool:
             return True
         if type(self) is not type(other):
             return False
+        getattr_ = getattr
         return all(
-            getattr(self, key) == getattr(other, key) for key in self.__immutable_keys__
+            getattr_(self, key) == getattr_(other, key) for key in self.__immutable_keys__
         )
 
     def __str__(self) -> str:
         fields = ", ".join(f"{_field_str(k, v)}" for k, v in self.__immutable_items__)
         return f"{type(self).__name__}({fields})"
 
     def __init__(self, **kwds: Any) -> None:
-        required: set[str] = set(self.__immutable_keys__)
-        if not required and not kwds:
-            # NOTE: Fastpath for empty slots
-            ...
-        elif required == set(kwds):
-            for name, value in kwds.items():
-                object.__setattr__(self, name, value)
-        elif missing := required.difference(kwds):
-            msg = (
-                f"{type(self).__name__!r} requires attributes {sorted(required)!r}, \n"
-                f"but missing values for {sorted(missing)!r}"
-            )
-            raise TypeError(msg)
-        else:
-            extra = set(kwds).difference(required)
-            msg = (
-                f"{type(self).__name__!r} only supports attributes {sorted(required)!r}, \n"
-                f"but got unknown arguments {sorted(extra)!r}"
-            )
-            raise TypeError(msg)
+        if (keys := self.__immutable_keys__) or kwds:
+            required = set(keys)
+            if required == kwds.keys():
+                object__setattr__ = object.__setattr__
+                for name, value in kwds.items():
+                    object__setattr__(self, name, value)
+            elif missing := required.difference(kwds):
+                raise _init_missing_error(self, required, missing)
+            else:
+                raise _init_extra_error(self, required, set(kwds).difference(required))
 
 
 def _field_str(name: str, value: Any) -> str:
@@ -149,3 +113,23 @@ def _field_str(name: str, value: Any) -> str:
     if isinstance(value, str):
         return f"{name}={value!r}"
     return f"{name}={value}"
+
+
+def _init_missing_error(
+    obj: object, required: Iterable[str], missing: Iterable[str]
+) -> TypeError:
+    msg = (
+        f"{type(obj).__name__!r} requires attributes {sorted(required)!r}, \n"
+        f"but missing values for {sorted(missing)!r}"
+    )
+    return TypeError(msg)
+
+
+def _init_extra_error(
+    obj: object, required: Iterable[str], extra: Iterable[str]
+) -> TypeError:
+    msg = (
+        f"{type(obj).__name__!r} only supports attributes {sorted(required)!r}, \n"
+        f"but got unknown arguments {sorted(extra)!r}"
+    )
+    return TypeError(msg)

narwhals/_plan/_meta.py

@@ -0,0 +1,91 @@
+"""Metaclasses and other unholy metaprogramming nonsense."""
+
+from __future__ import annotations
+
+# ruff: noqa: N806
+from itertools import chain
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any, Final, TypeVar
+
+    import _typeshed
+    from typing_extensions import dataclass_transform
+
+    from narwhals._plan.typing import Seq
+
+    T = TypeVar("T")
+
+else:
+    # https://docs.python.org/3/library/typing.html#typing.dataclass_transform
+    def dataclass_transform(
+        *,
+        eq_default: bool = True,
+        order_default: bool = False,
+        kw_only_default: bool = False,
+        frozen_default: bool = False,
+        field_specifiers: tuple[type[Any] | Callable[..., Any], ...] = (),
+        **kwargs: Any,
+    ) -> Callable[[T], T]:
+        def decorator(cls_or_fn: T) -> T:
+            cls_or_fn.__dataclass_transform__ = {
+                "eq_default": eq_default,
+                "order_default": order_default,
+                "kw_only_default": kw_only_default,
+                "frozen_default": frozen_default,
+                "field_specifiers": field_specifiers,
+                "kwargs": kwargs,
+            }
+            return cls_or_fn
+
+        return decorator
+
+
+__all__ = ["ImmutableMeta", "SlottedMeta", "dataclass_transform"]
+
+flatten = chain.from_iterable
+_KEYS_NAME: Final = "__immutable_keys__"
+_HASH_NAME: Final = "__immutable_hash_value__"
+
+
+class SlottedMeta(type):
+    """Ensure [`__slots__`] are always defined to prevent `__dict__` creation.
+
+    [`__slots__`]: https://docs.python.org/3/reference/datamodel.html#object.__slots__
+    """
+
+    # https://github.com/python/typeshed/blob/776508741d76b58f9dcb2aaf42f7d4596a48d580/stdlib/abc.pyi#L13-L19
+    # https://github.com/python/typeshed/blob/776508741d76b58f9dcb2aaf42f7d4596a48d580/stdlib/_typeshed/__init__.pyi#L36-L40
+    # https://github.com/astral-sh/ruff/issues/8353#issuecomment-1786238311
+    # https://docs.python.org/3/reference/datamodel.html#creating-the-class-object
+    def __new__(
+        metacls: type[_typeshed.Self],
+        cls_name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+        /,
+        **kwds: Any,
+    ) -> _typeshed.Self:
+        namespace.setdefault("__slots__", ())
+        return super().__new__(metacls, cls_name, bases, namespace, **kwds)  # type: ignore[no-any-return, misc]
+
+
+@dataclass_transform(kw_only_default=True, frozen_default=True)
+class ImmutableMeta(SlottedMeta):
+    def __new__(
+        metacls: type[_typeshed.Self],
+        cls_name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+        /,
+        **kwds: Any,
+    ) -> _typeshed.Self:
+        KEYS, HASH = _KEYS_NAME, _HASH_NAME
+        getattr_: Callable[..., Seq[str]] = getattr
+        it_bases = (getattr_(b, KEYS, ()) for b in bases)
+        it_all = chain(
+            flatten(it_bases), namespace.get(KEYS, namespace.get("__slots__", ()))
+        )
+        namespace[KEYS] = tuple(key for key in it_all if key != HASH)
+        return super().__new__(metacls, cls_name, bases, namespace, **kwds)  # type: ignore[no-any-return, misc]

narwhals/_plan/expr.py

@@ -194,13 +194,11 @@ def hist(
             if bin_count is not None:
                 msg = "can only provide one of `bin_count` or `bins`"
                 raise ComputeError(msg)
-            node = F.HistBins(bins=tuple(bins), include_breakpoint=include_breakpoint)
+            node = F.Hist.from_bins(bins, include_breakpoint=include_breakpoint)
         elif bin_count is not None:
-            node = F.HistBinCount(
-                bin_count=bin_count, include_breakpoint=include_breakpoint
-            )
+            node = F.Hist.from_bin_count(bin_count, include_breakpoint=include_breakpoint)
         else:
-            node = F.HistBinCount(include_breakpoint=include_breakpoint)
+            node = F.Hist.from_bin_count(include_breakpoint=include_breakpoint)
         return self._with_unary(node)
 
     def log(self, base: float = math.e) -> Self:

narwhals/_plan/expressions/__init__.py

@@ -10,9 +10,14 @@
 from narwhals._plan.expressions import (
     aggregation,
     boolean,
+    categorical,
     functions,
+    lists,
     operators,
     selectors,
+    strings,
+    struct,
+    temporal,
 )
 from narwhals._plan.expressions.aggregation import AggExpr, OrderableAggExpr, max, min
 from narwhals._plan.expressions.expr import (
@@ -85,10 +90,12 @@
     "_ColumnSelection",
     "aggregation",
     "boolean",
+    "categorical",
     "col",
     "cols",
     "functions",
     "index_columns",
+    "lists",
     "max",
     "min",
     "named_ir",
@@ -97,4 +104,7 @@
     "over",
     "over_ordered",
     "selectors",
+    "strings",
+    "struct",
+    "temporal",
 ]

narwhals/_plan/expressions/aggregation.py

@@ -50,18 +50,18 @@ class ArgMin(OrderableAggExpr): ...
 class ArgMax(OrderableAggExpr): ...
 # fmt: on
 class Quantile(AggExpr):
-    __slots__ = (*AggExpr.__slots__, "interpolation", "quantile")
+    __slots__ = ("interpolation", "quantile")
     quantile: float
     interpolation: RollingInterpolationMethod
 
 
 class Std(AggExpr):
-    __slots__ = (*AggExpr.__slots__, "ddof")
+    __slots__ = ("ddof",)
     ddof: int
 
 
 class Var(AggExpr):
-    __slots__ = (*AggExpr.__slots__, "ddof")
+    __slots__ = ("ddof",)
     ddof: int
 
 

narwhals/_plan/expressions/expr.py

@@ -407,7 +407,7 @@ class OrderedWindowExpr(
     child=("expr", "partition_by", "order_by"),
     config=ExprIROptions.renamed("over_ordered"),
 ):
-    __slots__ = ("expr", "partition_by", "order_by", "sort_options", "options")  # noqa: RUF023
+    __slots__ = ("order_by", "sort_options")
     expr: ExprIR
     partition_by: Seq[ExprIR]
     order_by: Seq[ExprIR]