Make a Python function subscriptable

for the purposes of a completely contrived programming exercise, I want to be able to access all objects that I have created of a specific class, using decorators, and by subscripting the class itself.

In the below example, I have created an instance of my A class, using the argument 1, and stored the result in a dictionary, using the key '1'.

So, by attempting to index the A class directly, it should return the instance of the object that was created, however, it only returns the parameter passed to __get__item, and raises a TypeError if I attempt to index it directly with A['1'].

import types


def remember(cl):
    seen = dict()

    def _remember(*args, **kwargs):
        key = ','.join(map(str, args))
        if key not in seen:
            seen[key] = cl(*args, **kwargs)
        return seen[key]

    _remember.__getitem__ = types.MethodType(seen.get, _remember)
    return _remember


@remember
class A(object):
    def __init__(self, x):
        pass


a = A(1)
b = A.__getitem__('1')
print(a == b) #should print true, but instead, A.__get__item returns '1'. Ideally, A['1'] would work.

Solution

Here's what I came up with.

We can't have A[<something>] as possible usage because that would mean changing the attributes of type. I've created A.get(<something>) instead for that use case.

This differs from your attempt in that rather than returning a function from the decorator, we return a new class that extends from the original. That class stores all instances of instantiations of its parent class, and indexes them by a key which is created from the args and kwargs.

The __getitem__ has been overridden on the class, and then in the parent class, and it takes any of:

A tuple of args (tuple) and kwargs (dict)
Args as a tuple
A single arg

Anyway, on with the code:

from typing import Any, Dict


def remember(klass):
    def get_key(args, kwargs):
        return hash((args, tuple(sorted(kwargs.items()))))

    class RememberedClass(klass):
        _instances: Dict[int, Any] = {}

        def __new__(cls, *args, **kwargs):
            c = klass(*args, **kwargs)
            klass.__getitem__ = cls.__getitem__

            key = get_key(args, kwargs)
            cls._instances[key] = c
            return c

        @staticmethod
        def get(*args, **kwargs) -> klass:
            key = get_key(args, kwargs)
            return RememberedClass._instances[key]

        def __getitem__(self, item) -> klass:
            if isinstance(item, tuple) and len(item) == 2:
                # (args, kwargs)
                args, kwargs = item
                key = get_key(args, kwargs)
            elif isinstance(item, tuple):
                # (arg, arg, ...)
                key = get_key(item, {})
            else:
                # arg
                key = get_key((item,), {})
            return RememberedClass._instances.get(key)

    return RememberedClass


@remember
class Foo:
    def __init__(self, val):
        pass

I believe this is what you wanted?

>>> f = Foo(1)

>>> g = Foo(2)

>>> f is g
False

>>> f is f[1]
True

>>> g is g[2]
True

>>> g is f[2]
True

>>> f is Foo.get(1)
True

>>> type(f)
__main__.Foo

>>> Foo
__main__.remember.<locals>.RememberedClass