From 02a7cd86c1187e1e43e116b34da2993871474057 Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 15 Oct 2026 20:42:04 -0400 Subject: [PATCH 1/4] =?UTF-8?q?=F0=9F=90=9B=20fix:=20keep=20the=20frozendi?= =?UTF-8?q?ct=20backport's=20storage=20out=20of=20subclasses'=20reach?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backport kept its contents in a slot named `_data`. The built-in has no such attribute, so on Python 3.15 a subclass is free to use `_data` for itself; on the backport, assigning it rebound the slot and replaced the mapping's contents. galax and coordinax both declared a `_data` field on their subclasses. Name-mangle the slot to `_frozendict__data`, so `self._data` on a subclass is just another attribute, as with the built-in. Pickles are unaffected: they rebuild through `__new__` from the data, not from slot state. Co-authored-by: Claude Opus 5.5 --- src/xmmutablemap/_frozendict.py | 46 +++++++++++++++++---------------- tests/test_immutablemap.py | 16 ++++++++++++ 2 files changed, 40 insertions(+), 22 deletions(-) diff --git a/src/xmmutablemap/_frozendict.py b/src/xmmutablemap/_frozendict.py index dd9fa5d..c9a4b60 100644 --- a/src/xmmutablemap/_frozendict.py +++ b/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,7 @@ 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] + self.__data = dict(*args, **kwargs) # type: ignore[assignment] return self # =========================================== @@ -110,7 +112,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 +124,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 +136,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 +148,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 +160,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 +172,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 +184,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 +196,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 +218,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 +238,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 +274,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 +335,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 +365,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 +395,9 @@ def __eq__(self, other: object) -> bool: """ if isinstance(other, frozendict): - return self._data == other._data + 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 +416,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 +437,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}()" diff --git a/tests/test_immutablemap.py b/tests/test_immutablemap.py index d61f97c..2263e78 100644 --- a/tests/test_immutablemap.py +++ b/tests/test_immutablemap.py @@ -356,6 +356,22 @@ 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_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) From fb5cf154b28e3800737fd2d4cf92ea1526f195bf Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 15 Oct 2026 20:47:39 -0400 Subject: [PATCH 2/4] =?UTF-8?q?=F0=9F=90=9B=20fix:=20build=20the=20frozend?= =?UTF-8?q?ict=20backport=20without=20a=20subclass's=20`=5F=5Fsetattr=5F?= =?UTF-8?q?=5F`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The built-in fills itself in C, so a subclass that overrides `__setattr__` constructs normally: a frozen dataclass subclass works on Python 3.15. The backport assigned its storage with `self.__data = ...`, which goes through the subclass's `__setattr__` and raised `FrozenInstanceError` there instead. Assign it with `object.__setattr__`. This also makes an `equinox.Module` subclass behave the same on every Python: its `__new__` drops the arguments, so it is an empty mapping, as on 3.15, rather than failing only on older Pythons. Co-authored-by: Claude Opus 5.5 --- src/xmmutablemap/_frozendict.py | 4 +++- tests/test_immutablemap.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/xmmutablemap/_frozendict.py b/src/xmmutablemap/_frozendict.py index c9a4b60..d1e489d 100644 --- a/src/xmmutablemap/_frozendict.py +++ b/src/xmmutablemap/_frozendict.py @@ -96,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 # =========================================== diff --git a/tests/test_immutablemap.py b/tests/test_immutablemap.py index 2263e78..d5cd4e0 100644 --- a/tests/test_immutablemap.py +++ b/tests/test_immutablemap.py @@ -372,6 +372,19 @@ def __init__(self, *_args: Any, **_kwargs: Any) -> None: 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) From ed949447aae9ec645790d2565eb9c477df443753 Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 15 Oct 2026 20:48:04 -0400 Subject: [PATCH 3/4] =?UTF-8?q?=F0=9F=93=9D=20docs:=20document=20subclassi?= =?UTF-8?q?ng,=20and=20fix=20the=20`tree=5Funflatten`=20example?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since 0.3.0 `ImmutableMap` is a `frozendict`, which takes its contents in `__new__`. Say so in a README section on subclassing: override `__new__` to transform input, register the subclass as a PyTree node, and don't combine it with an `equinox.Module` or a dataclass, which galax and coordinax both did (via `ImmutableMap.__init__` and the private `_data`, now gone). The `tree_unflatten` doctest passed `tree_flatten`'s `(children, keys)` as `(aux_data, children)`, and documented the swapped result `ImmutableMap({1: 'a', 2: 'b'})`; the `jax.tree.unflatten` example passed that pair as the leaves. Both now round-trip `d`. Co-authored-by: Claude Opus 5.5 --- README.md | 35 +++++++++++++++++++++++++++++++++++ src/xmmutablemap/_core.py | 11 ++++++----- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 251c84b..61630e0 100644 --- a/README.md +++ b/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 diff --git a/src/xmmutablemap/_core.py b/src/xmmutablemap/_core.py index 03c4aaf..ca0a0d9 100644 --- a/src/xmmutablemap/_core.py +++ b/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))) From a2451a6e6caf4d7766acca5473c8e9ddfebf951d Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 15 Oct 2026 21:26:01 -0400 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=9A=A8=20fix-lint:=20allow=20`frozend?= =?UTF-8?q?ict.=5F=5Feq=5F=5F`=20to=20read=20the=20other's=20storage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pylint's protected-access fires on `other.__data` now that the slot is name-mangled; `other` is a `frozendict` too. Co-authored-by: Claude Opus 5.5 --- src/xmmutablemap/_frozendict.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/xmmutablemap/_frozendict.py b/src/xmmutablemap/_frozendict.py index d1e489d..ab94695 100644 --- a/src/xmmutablemap/_frozendict.py +++ b/src/xmmutablemap/_frozendict.py @@ -397,6 +397,7 @@ def __eq__(self, other: object) -> bool: """ if isinstance(other, frozendict): + # pylint: disable-next=protected-access # same class return self.__data == other.__data if isinstance(other, Mapping): return self.__data == dict(other)