Python WeakKeyDictionary for unhashable types

Question

As raised in cpython issue 88306, python WeakKeyDictionary fails for non hashable types. According to the discussion in the python issue above, this is an unnecessary restriction, using ids of the keys instead of hash would work just fine: In this special case ids are unique identifiers for the keys in the WeakKeyDictionary, because the keys are automatically removed when the original object is deleted. It is important to be aware that using ids instead of hashes is only feasible in this very special case.

We can tweak weakref.WeakKeyDictionary (see gist) to achieve the desired behaviour. In summary, this implementation wraps the weakref keys as follows:

class _IdKey:
    def __init__(self, key):
        self._id = id(key)

    def __hash__(self):
        return self._id

    def __eq__(self, other: typing_extensions.Self):
        return self._id == other._id

    def __repr__(self):
        return f"<_IdKey(_id={self._id})>"


class _IdWeakRef(_IdKey):
    def __init__(self, key, remove: typing.Callable[[typing.Any], None]):
        super().__init__(key)
        # hold weak ref to avoid garbage collection of the remove callback
        self._ref = weakref.ref(key, lambda _: remove(self))

    def __call__(self):
        # used in weakref.WeakKeyDictionary.__copy__
        return self._ref()

    def __repr__(self):
        return f"<_IdKey(_id={self._id},{self._ref})>"

class WeakKeyIdDictionary(weakref.WeakKeyDictionary):
   """
   overrides all methods involving dictionary access key 
   """
   ... https://gist.github.com/barmettl/b198f0cf6c22047df77483e8aa28f408

However, this depends on the details of the implementation of weakref.WeakKeyDictionary (using python3.10 here) and is likely to break in future (or even past) versions of python. Of course, alternatively one can just rewrite an entirely new class.

It is also possible to implement a custom __hash__ method for all classes, but this won't work when dealing with external code and will give unreliable hashes for use cases beyond weakref.WeakKeyDictionary. We can also monkey patch __hash__, but this is not possible in particular for built in classes and will have unintended effects in other parts of the code.

Thus the following question: How should one store non hashable items in a WeakKeyDictionary?

Answer 1

There is a way which does not rely on knowing the internals of WeakKeyDictionary:

from weakref import WeakKeyDictionary, WeakValueDictionary

class Id:
    def __init__(self, key):
        self._id = id(key)
    def __hash__(self):
        return self._id
    def __eq__(self, other):
        return self._id == other._id

class WeakUnhashableKeyDictionary:
    def __init__(self, *args, **kwargs):
        # TODO Do something to initialize given args and kwargs.
        self.keys = WeakValueDictionary()
        self.values = WeakKeyDictionary()

    def __getitem__(self, key):
        return self.values.__getitem__(Id(key))
    
    def __setitem__(self, key, value):
        _id = Id(key)
        # NOTE This works because key holds on _id iif key exists,
        # and _id holds on value iif _id exists. Transitivity. QED.
        # Because key is only stored as a value, it does not need to be hashable.
        self.keys.__setitem__(_id, key)
        self.values.__setitem__(_id, value)

    def __delitem__(self, key):
        self.keys.__delitem__(Id(key))
        self.values.__delitem__(Id(key))
        
    # etc. other methods should be relatively simple to implement.
    # TODO Might require some locks or care in the ordering of operations to work threaded.
    # TODO Add clean error handling.

This is just a generalization of my answer to a method caching problem.