Utility decorator classes (subscriptable, typedmethod) that allow you to add subscript notation (__getitem__ hooks) to class methods

example_1.py

from subscriptable import subscriptable

@subscriptable
def other_typed_value(value, **kwargs):
    subscript = subscriptable.key(kwargs)
    print(subscript, value)

value = other_typed_value[int]('12345') # <class 'int'> 12345
value = other_typed_value('12345') # None 12345
print("Standard function str:", other_typed_value) # Standard function str: <subscriptable <function other_typed_value at 0x102d6e560>>

example_2.py

from subscriptable import subscriptable

@subscriptable.container
class MyClass:

    @subscriptable
    def compute_typed_value(self, value, *args, **kwargs):
        print(self, args, kwargs)
        if subscriptable.has_key(kwargs):
            value = subscriptable.key(kwargs)(value)
        self.my_other_method()
        return value

    def my_other_method(self):
        print('Doing some other things!')
        return 3

a = MyClass() 
value = a.compute_typed_value[int]('12345') # <__main__.MyClass object at 0x104ac3a60> () {'___SUBSCRIPT_KEY': <class 'int'>}
print(value, type(value)) # 12345 <class 'int'>
value = a.compute_typed_value('12345') # <__main__.MyClass object at 0x104ac3a60> () {}
print(value, type(value)) # 12345 <class 'str'>
print("Class Method str:", a.compute_typed_value) # Class Method str: <subscriptable <bound method MyClass.compute_typed_value of <__main__.MyClass object at 0x104ac3a60>>>

example_3.py

from subscriptable import subscriptable, typedmethod
@subscriptable.container
class MyClass2:

    @typedmethod
    def compute_typed_value(self, value, *args, **kwargs):
        print(self, args, kwargs)
        self.my_other_method()
        return value

    def my_other_method(self):
        print('Doing some other things!')
        return 3

a = MyClass2()
value = a.compute_typed_value[int]('12345') # <__main__.MyClass2 object at 0x10b58fa60> () {}
print(value, type(value)) # 12345 <class 'int'>
value = a.compute_typed_value('12345') # <__main__.MyClass2 object at 0x10b58fa60> () {}
print(value, type(value)) # 12345 <class 'str'>
print("Class Method str:", a.compute_typed_value) # Class Method str: <typedmethod <bound method MyClass2.compute_typed_value of <__main__.MyClass2 object at 0x10b58fa60>>>

subscriptable.py

from functools import partial, wraps
from types import MethodType, FunctionType
from typing import Any, Callable

class subscriptable:
    _SUBSCRIPT_KEY = '___SUBSCRIPT_KEY'
    _HAS_SUBSCRIPTABLE_METHODS = '___HAS_SUBSCRIPTABLE_METHODS'

    _callout: Callable
    def __init__(self, callout: Callable, instance: Any = None):
        if instance is not None and isinstance(callout, FunctionType):
            self._callout = MethodType(callout, instance)
        else:
            self._callout = callout

    def bind(self, instance):
        return self.__class__(self._callout, instance=instance)

    def __getitem__(self, specified_type):
        return partial(self.__call__, **{self.__class__._SUBSCRIPT_KEY: specified_type})
    
    def __call__(self, *args, **kwargs):
        """A transparent passthrough to the wrapped method"""
        return self._callout(*args, **kwargs)
    
    def __str__(self):
        return f"<{self.__class__.__name__} {self._callout}>"
    
    @classmethod
    def has_key(cls, foo_kwargs):
        """A utility method to determine whether the provided kwargs has the expected subscript key"""
        return cls._SUBSCRIPT_KEY in foo_kwargs
    
    @classmethod
    def key(cls, foo_kwargs:dict):
        """A utility method that allows the subscript key to be consumed by the wrapped method, without needing to know the inner workings"""
        return foo_kwargs.pop(cls._SUBSCRIPT_KEY, None)
    
    @classmethod
    def container(cls, clazz):
        """A decorator for classes containing `subscriptable` methods"""
        if not hasattr(clazz, cls._HAS_SUBSCRIPTABLE_METHODS):
            orig_init = clazz.__init__
            @wraps(clazz.__init__)
            def __init__(self, *args, **kwargs):
                for attr_name in dir(self):
                    attr_value = getattr(self, attr_name)
                    if isinstance(attr_value, cls):
                        setattr(self, attr_name, attr_value.bind(self))
                orig_init(self, *args, **kwargs)
            clazz.__init__ = __init__
            setattr(clazz, cls._HAS_SUBSCRIPTABLE_METHODS, True)
        return clazz
    

class typedmethod(subscriptable):

    def __init__(self, *args, **kwargs):
        self._spread_method = kwargs.pop('splat', None)
        super().__init__(*args, **kwargs)

    def __call__(self, *args, **kwargs):
        specified_type = subscriptable.key(kwargs)
        print(specified_type)
        if specified_type is not None and not isinstance(specified_type, type):
            raise TypeError('Subscript type must be a type')
        result = super().__call__(*args, **kwargs)
        if specified_type is not None:
            if self._spread_method == '*':
                result = specified_type(*result)
            elif self._spread_method == '**':
                result = specified_type(**result)
            else:
                result = specified_type(result)
        return result
    
    @classmethod
    def config(cls, **config_values):
        return partial(cls, **config_values)