🐛 fix: build the frozendict backport the way the built-in does; document subclassing

README.md

@@ -120,6 +120,41 @@ while `frozendict | other` merges only a `dict` or `frozendict` (so, for
 example, a custom `Mapping` with no `|` of its own raises `TypeError`). Convert
 such mappings with `dict(...)` when switching to the built-in.
 
+#### Subclassing
+
+Being a `frozendict`, an `ImmutableMap` gets its contents in `__new__`, and they
+cannot change afterwards. A subclass that wants to transform its input overrides
+`__new__`; its `__init__`, if any, runs on an already-filled map. Subclasses are
+not PyTree nodes unless they register themselves:
+
+```python
+from jax.tree_util import register_pytree_node_class
+
+
+@register_pytree_node_class
+class UpperKeys(ImmutableMap):
+    def __new__(cls, *args, **kwargs):
+        data = dict(*args, **kwargs)
+        return super().__new__(cls, {k.upper(): v for k, v in data.items()})
+
+
+print(UpperKeys(a=1))
+# UpperKeys({'A': 1})
+print(jax.tree.leaves(UpperKeys(a=1, b=2)))
+# [1, 2]
+```
+
+> [!WARNING]
+>
+> **Don't subclass `ImmutableMap` with an `equinox.Module` or a dataclass.** The
+> contents are whatever `__new__` receives, as `dict(*args, **kwargs)`, and
+> nothing in `__init__` can change them. An `equinox.Module` drops its arguments
+> in `__new__`, so it is always an empty mapping; a dataclass gets its
+> constructor arguments as its contents, not its fields. Before 0.3.0 this
+> pattern worked by calling `ImmutableMap.__init__` and setting the private
+> `_data`; neither exists any more. Have such a class subclass
+> `collections.abc.Mapping` and keep its own `dict` instead.
+
 ### JAX Integration
 
 One of the key benefits of `ImmutableMap` is its compatibility with JAX. Since

src/xmmutablemap/_core.py

@@ -175,12 +175,13 @@ def tree_unflatten(
         >>> import jax
         >>> from xmmutablemap import ImmutableMap
         >>> d = ImmutableMap(a=1, b=2)
-        >>> flat = d.tree_flatten()
-        >>> ImmutableMap.tree_unflatten(*flat)
-        ImmutableMap({1: 'a', 2: 'b'})
+        >>> children, keys = d.tree_flatten()
+        >>> ImmutableMap.tree_unflatten(keys, children)
+        ImmutableMap({'a': 1, 'b': 2})
 
-        >>> jax.tree.unflatten(jax.tree.structure(d), flat)
-        ImmutableMap({'a': (1, 2), 'b': ('a', 'b')})
+        >>> leaves, treedef = jax.tree.flatten(d)
+        >>> jax.tree.unflatten(treedef, leaves)
+        ImmutableMap({'a': 1, 'b': 2})
 
         """
         return cls(tuple(zip(aux_data, children, strict=True)))

src/xmmutablemap/_frozendict.py

@@ -72,10 +72,12 @@ class frozendict(Mapping[K, V]):  # noqa: N801  # pylint: disable=invalid-name
     # `__slots__` alone gives instances no `__dict__`, so attributes can be
     # neither set nor deleted -- with the same message the built-in raises.
     # No `__setattr__` override is needed to reproduce that.
-    # ponytail: `_data` itself stays rebindable, being a slot; it is private.
-    __slots__ = ("_data",)
+    # The storage is name-mangled (to `_frozendict__data`): the built-in has no
+    # such attribute, so a subclass that sets its own `_data` must not reach it.
+    # ponytail: it stays rebindable, being a slot; it is private.
+    __slots__ = ("__data",)
 
-    _data: dict[K, V]
+    __data: dict[K, V]
 
     # Built in `__new__`, not `__init__`, as the built-in is: calling
     # `__init__` again must not change the contents, and a subclass's own
@@ -94,7 +96,9 @@ def __new__(
         if cls is frozendict and args and not kwargs and type(args[0]) is frozendict:
             return cast("Self", args[0])
         self = super().__new__(cls)
-        self._data = dict(*args, **kwargs)  # type: ignore[assignment]
+        # Not `self.__data = ...`: the built-in fills itself in C, so a
+        # subclass's `__setattr__` (e.g. a frozen dataclass's) never sees it.
+        object.__setattr__(self, "_frozendict__data", dict(*args, **kwargs))
         return self
 
     # ===========================================
@@ -110,7 +114,7 @@ def __getitem__(self, key: K) -> V:
         1
 
         """
-        return self._data[key]
+        return self.__data[key]
 
     def __iter__(self) -> Iterator[K]:
         """Iterate over the keys, in insertion order.
@@ -122,7 +126,7 @@ def __iter__(self) -> Iterator[K]:
         ['a', 'b']
 
         """
-        return iter(self._data)
+        return iter(self.__data)
 
     def __len__(self) -> int:
         """Return the number of items.
@@ -134,7 +138,7 @@ def __len__(self) -> int:
         2
 
         """
-        return len(self._data)
+        return len(self.__data)
 
     def __contains__(self, key: object) -> bool:
         """Check whether a key is present.
@@ -146,7 +150,7 @@ def __contains__(self, key: object) -> bool:
         True
 
         """
-        return key in self._data
+        return key in self.__data
 
     def __reversed__(self) -> Iterator[K]:
         """Iterate over the keys in reverse insertion order.
@@ -158,7 +162,7 @@ def __reversed__(self) -> Iterator[K]:
         ['b', 'a']
 
         """
-        return reversed(self._data)
+        return reversed(self.__data)
 
     def keys(self) -> KeysView[K]:
         """Return a view of the keys.
@@ -170,7 +174,7 @@ def keys(self) -> KeysView[K]:
         dict_keys(['a', 'b'])
 
         """
-        return self._data.keys()
+        return self.__data.keys()
 
     def values(self) -> ValuesView[V]:
         """Return a view of the values.
@@ -182,7 +186,7 @@ def values(self) -> ValuesView[V]:
         dict_values([1, 2])
 
         """
-        return self._data.values()
+        return self.__data.values()
 
     def items(self) -> ItemsView[K, V]:
         """Return a view of the items.
@@ -194,7 +198,7 @@ def items(self) -> ItemsView[K, V]:
         dict_items([('a', 1), ('b', 2)])
 
         """
-        return self._data.items()
+        return self.__data.items()
 
     @overload
     def get(self, key: K, /) -> V | None: ...
@@ -216,7 +220,7 @@ def get(self, key: K, /, default: V | _T | None = None) -> V | _T | None:
         2
 
         """
-        return self._data.get(key, default)
+        return self.__data.get(key, default)
 
     def copy(self, /) -> "frozendict[K, V]":
         """Return a shallow copy.
@@ -236,7 +240,7 @@ def copy(self, /) -> "frozendict[K, V]":
         # `PyFrozenDict_CheckExact` here, so a *subclass* must not return
         # `self`. `ImmutableMap` overrides `copy` precisely because of this.
         # pylint: disable-next=unidiomatic-typecheck
-        return self if type(self) is frozendict else frozendict(self._data)
+        return self if type(self) is frozendict else frozendict(self.__data)
 
     def __copy__(self) -> "frozendict[K, V]":
         """Support `copy.copy`.
@@ -272,7 +276,7 @@ def __getnewargs__(self) -> tuple[dict[K, V]]:
         ({'a': 1},)
 
         """
-        return (dict(self._data),)
+        return (dict(self.__data),)
 
     def __reduce_ex__(self, protocol: SupportsIndex, /) -> tuple[Any, ...]:
         """Reduce for pickling and `copy`, exactly as the built-in does.
@@ -333,7 +337,7 @@ def __or__(self, other: Any, /) -> "frozendict[Any, Any]":
             return NotImplemented
         # The built-in's C implementation hardcodes the result type, so a
         # subclass of `frozendict` gets a plain `frozendict` back. Match that.
-        return frozendict(self._data | dict(other))
+        return frozendict(self.__data | dict(other))
 
     def __ror__(self, other: Any, /) -> Any:
         """Merge into another mapping.
@@ -363,7 +367,7 @@ def __ror__(self, other: Any, /) -> Any:
             return NotImplemented
         # Let the dict drive the merge so that e.g. `OrderedDict | frozendict`
         # stays an OrderedDict.
-        return other | self._data
+        return other | self.__data
 
     # ===========================================
     # Other
@@ -393,9 +397,10 @@ def __eq__(self, other: object) -> bool:
 
         """
         if isinstance(other, frozendict):
-            return self._data == other._data
+            # pylint: disable-next=protected-access  # same class
+            return self.__data == other.__data
         if isinstance(other, Mapping):
-            return self._data == dict(other)
+            return self.__data == dict(other)
         return NotImplemented
 
     def __hash__(self) -> int:
@@ -414,11 +419,11 @@ def __hash__(self) -> int:
 
         """
         try:
-            return hash(frozenset(self._data.items()))
+            return hash(frozenset(self.__data.items()))
         except TypeError:
             # Re-raise as the built-in does, naming the unhashable value
             # rather than the `(key, value)` tuple holding it.
-            for value in self._data.values():
+            for value in self.__data.values():
                 hash(value)
             raise  # pragma: no cover  # only if a value's hash is flaky
 
@@ -435,4 +440,4 @@ def __repr__(self) -> str:
 
         """
         name = type(self).__name__
-        return f"{name}({self._data!r})" if self._data else f"{name}()"
+        return f"{name}({self.__data!r})" if self.__data else f"{name}()"

tests/test_immutablemap.py

@@ -356,6 +356,35 @@ def __init__(self, *_args: Any, **_kwargs: Any) -> None:
 
         assert Sub(a=1).seen == {"a": 1}
 
+    def test_subclass_data_attribute_is_its_own(self) -> None:
+        """A subclass attribute named `_data` does not touch the contents.
+
+        The built-in has no `_data`, so there it is just another attribute;
+        the backport must not keep its storage under a name a subclass can
+        reach by accident.
+        """
+
+        class Sub(frozendict[str, int]):
+            def __init__(self, *_args: Any, **_kwargs: Any) -> None:
+                self._data = {"z": 0}
+
+        s = Sub(a=1)
+        assert dict(s) == {"a": 1}
+        assert s._data == {"z": 0}
+
+    def test_construction_bypasses_subclass_setattr(self) -> None:
+        """A subclass `__setattr__` is not involved in building the contents.
+
+        The built-in fills itself in C, so e.g. a frozen dataclass subclass
+        constructs; the backport must not go through `__setattr__` either.
+        """
+
+        class NoSetattr(frozendict[str, int]):
+            def __setattr__(self, name: str, value: Any) -> None:
+                raise AttributeError(name)
+
+        assert dict(NoSetattr(a=1)) == {"a": 1}
+
     def test_hash_is_frozenset_of_items(self) -> None:
         """The hash is that of the items' frozenset, as for the built-in."""
         d = frozendict(a=1, b=2)