Parameterized decorator that works on functions, and methods (instance, class, or static).

deco.py

#!/usr/bin/env python3

'''
$ ./deco.py
Test methods...
self.rstate: type=<class 'str'>, value=asdf
self.rstate: type=<class 'int'>, value=123
Unknown rstate.payload: <class 'int'>
Unknown rstate.payload: <class 'str'>
Unknown rstate.payload: <class 'int'>

Test classmethods...
cls.rstate_str: type=<class 'str'>, value=asdf
cls.rstate_str: type=<class 'str'>, value=asdf
cls.rstate_int: type=<class 'int'>, value=123
Unknown rstate.payload: <class 'int'>

Test staticmethods...
decorator -> staticmethod: sta.rstate_str: type=<class 'str'>, value=asdf
decorator -> staticmethod: sta.rstate_str: type=<class 'str'>, value=asdf
staticmethod -> decorator: sta.rstate_int: type=<class 'int'>, value=123
staticmethod -> decorator: sta.rstate_int: type=<class 'int'>, value=123
Unknown rstate.payload: <class 'int'>

Test functions...
f_rstate_str: type=<class 'str'>, value=asdf
f_rstate_int: type=<class 'int'>, value=123
Unknown rstate.payload: <class 'int'>
Unknown rstate.payload: <class 'str'>
'''

from collections.abc import Sequence

class allow_payload(object):
    '''Given a decorated function and its argument at a specified position,
    ensure that argument.payload is an instance of specified types.
    '''

    def __init__(self, klasses=(), arg_pos=0):
        if isinstance(klasses, type):
            self.klasses=(klasses,)
        elif isinstance(klasses, Sequence):
            self.klasses = tuple(klasses)
        else:
            raise TypeError(f'{klasses} is not instance of {{type, Sequence}}')

        if arg_pos < 0:
            raise ValueError(f'Invalid arg_pos: {arg_pos}')
        self.arg_pos = arg_pos

        self.f = None

    def __set_name__(self, owner, name):
        self.name = name

    def __call__(self, *args, **kwargs):
        if self.f is None:
            # During decorator definition
            self.f = args[0]
            return self
        else:
            # Decorated function is called
            decorated_function = self.wrap(None)
            return decorated_function(*args, **kwargs)

    def __get__(self, instance, cls):
        # Decorated method is called
        if instance is None:
            decorated_method = self.wrap(cls)
        else:
            decorated_method = self.wrap(instance)
        return decorated_method

    def wrap(self, owner):
        '''
        :param owner: instance or cls
        '''
        def wrap2(*args, **kwargs):
            '''Check payload's type.'''
            if owner is None or isinstance(self.f, staticmethod):
                # Plain old function or static method
                rstate = args[self.arg_pos]
            else:
                if self.arg_pos > 0:
                    rstate = args[self.arg_pos-1]
                else:
                    rstate = owner

            payload = rstate.payload

            if not isinstance(payload, self.klasses):
                raise TypeError(f'Unknown rstate.payload: {payload.__class__}')

            if owner is None:
                # Plain old function
                return self.f(*args, **kwargs)

            # When instance methods are more common than class/static method,
            # test for instance methods first for optimization.
            if callable(self.f):
                return self.f(owner, *args, **kwargs)

            if isinstance(self.f, staticmethod):
                return self.f.__func__(*args, **kwargs)

            if isinstance(self.f, classmethod):
                return self.f.__func__(owner, *args, **kwargs)

            raise TypeError('FATAL BUG. CONTACT MAINTAINER ASAP!')
        return wrap2

class RunState(object):
    def __init__(self, payload):
        self.payload = payload

from abc import ABC, abstractmethod
class TestABC(ABC):
    @allow_payload(str, 1)
    @abstractmethod
    def rstate_str(self, rstate):
        pass

    @allow_payload(int, 1)
    @abstractmethod
    def rstate_int(self, rstate):
        pass

class Test(TestABC):
    def __init__(self):
        pass

    @allow_payload(str, 1)
    def rstate_str(self, rstate):
        print(f'self.rstate: type={type(rstate.payload)}, value={rstate.payload}')

    @allow_payload(int, 1)
    def rstate_int(self, rstate):
        print(f'self.rstate: type={type(rstate.payload)}, value={rstate.payload}')

    @allow_payload(str, 1)
    def rstate_str_liar(self, _):
        rstate_liar = RunState(98765)
        super().rstate_str(rstate_liar)
        print(f'self.rstate liar: type={type(rstate.payload)}, value={rstate.payload}')

    @allow_payload(str, 1)
    @classmethod
    def c_rstate_str(cls, rstate):
        print(f'cls.rstate_str: type={type(rstate.payload)}, value={rstate.payload}')

    @classmethod
    @allow_payload(int, 1)
    def c_rstate_int(cls, rstate):
        print(f'cls.rstate_int: type={type(rstate.payload)}, value={rstate.payload}')

    @allow_payload(str)
    @staticmethod
    def s_rstate_str(rstate):
        print(f'sta.rstate_str: type={type(rstate.payload)}, value={rstate.payload}')

    @staticmethod
    @allow_payload(int)
    def s_rstate_int(rstate):
        print(f'sta.rstate_int: type={type(rstate.payload)}, value={rstate.payload}')

@allow_payload(str)
def f_rstate_str(rstate):
    print(f'f_rstate_str: type={type(rstate.payload)}, value={rstate.payload}')

@allow_payload(int)
def f_rstate_int(rstate):
    print(f'f_rstate_int: type={type(rstate.payload)}, value={rstate.payload}')

rstate_str = RunState('asdf')
rstate_int = RunState(123)

print('Test methods...')
t = Test()
t.rstate_str(rstate_str)
t.rstate_int(rstate_int)
try: t.rstate_str(rstate_int)
except Exception as e: print(e)
try: t.rstate_int(rstate_str)
except Exception as e: print(e)
try: t.rstate_str_liar(rstate_str)
except Exception as e: print(e)
print()

print('Test classmethods...')
Test.c_rstate_str(rstate_str)
t.c_rstate_str(rstate_str)
Test.c_rstate_int(rstate_int)
try: Test.c_rstate_str(rstate_int)
except Exception as e: print(e)
print()

print('Test staticmethods...')
print('decorator -> staticmethod:', end=' '); Test.s_rstate_str(rstate_str)
print('decorator -> staticmethod:', end=' '); t.s_rstate_str(rstate_str)
print('staticmethod -> decorator:', end=' '); Test.s_rstate_int(rstate_int)
print('staticmethod -> decorator:', end=' '); t.s_rstate_int(rstate_int)
try: Test.s_rstate_str(rstate_int)
except Exception as e: print(e)
print()

print('Test functions...')
f_rstate_str(rstate_str)
f_rstate_int(rstate_int)
try: f_rstate_str(rstate_int)
except Exception as e: print(e)
try: f_rstate_int(rstate_str)
except Exception as e: print(e)