Search code examples
pythonpython-decorators

Best way to make decorator for unique class instances Python


I am working with Python code that calls into C wrappers, but the C code is very buggy (and I have no other alternative) and causes a segfault when a Python object managed in C goes out of scope, so I have to keep references to each object created.

Is there any good way to make an ergonomic "unique" wrapper where each class can only have one instance per set of constructor arguments, e.g.

@unique
class Test:
    cls_val = 0
    def __init__(self, val):
        self.val = val

a = Test(1)
b = Test(1)

assert a is b

c = Test(2)
d = Test(2)

assert c is not b and c is not a
assert c is d

I've made this decorator, but it prevents any @unique-decorated class from being used as a base class (constructing an instance of the derived class calls the __new__ of the decorator).

def unique(unique_cls):
    class Unique:
        instances = {}
        unique_class = unique_cls
        def __new__(cls, *args, **kwargs):
            if not Unique.instances.get(
                (
                     cls.unique_class,
                     f_args := frozenset(args),
                     f_kwargs := frozenset(kwargs),
                )                
            ):
                Unique.instances[
                     (Unique.unique_class, f_args, f_kwargs)
                ] = Unique.unique_class(*args, **kwargs)
            return Unique.instances[(Unique.unique_class, f_args, f_kwargs)]
        def __getattr__(self, name):
            # Overloaded to get class attributes working for decorated classes
            return object.__getattribute__(Unique.unique_class, name)
    return Unique

Solution

  • You can use functools.lru_cache to store all instances in a cache, and set no limit on the cache size. When you call the constructor, you'll get a new instance, and it will be stored in the cache. Then whenever you call the constructor with the same arguments again, you'll get the cached instance. This also means that every object always has a reference from the cache.

    from functools import lru_cache
    
    @lru_cache(maxsize=None)
    class Test:
        def __init__(self, val):
            self.val = val
    

    Demonstration:

    >>> a = Test(1)
    >>> a2 = Test(1)
    >>> a is a2
    True
    >>> b = Test(2)
    >>> b2 = Test(2)
    >>> b is b2
    True
    >>> a is b
    False
    

    If you need to be able to subclass Test (or really, do anything with Test itself except create instances), then you can override __new__ and apply the decorator there. This works because cls is an argument to __new__, so the cache will distinguish between different instances by their class as well as by their __init__ arguments.

    class Test:
        @lru_cache(maxsize=None)
        def __new__(cls, *args, **kwargs):
            return object.__new__(cls)
        def __init__(self, val):
            self.val = val
    

    Demonstration:

    >>> Test(1) is Test(1)
    True
    >>> Test(1) is Test(2)
    False
    >>> class SubTest(Test): pass
    ... 
    >>> Test(1) is SubTest(1)
    False
    >>> SubTest(1) is SubTest(1)
    True