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))) diff --git a/src/xmmutablemap/_frozendict.py b/src/xmmutablemap/_frozendict.py index dd9fa5d..ab94695 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,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}()" diff --git a/tests/test_immutablemap.py b/tests/test_immutablemap.py index d61f97c..d5cd4e0 100644 --- a/tests/test_immutablemap.py +++ b/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)