CODEWITHSUNDEEP
pythonpython-decorators
redhotsnow
redhotsnow

Reputation: 95

How to check argument and return types via a decorator in Python

I have two simple decorators, which do basic type-check of function argument and return types. They work separately, but not in combination. And yes, I am aware of the static type checking in Python 3.6, but I still want to get this working :)

I have adapted the Python 2 code from the example here: https://www.python.org/dev/peps/pep-0318/#examples

from functools import wraps

def accepts(*types):
    def check_accepts(f):
        assert len(types) == f.__code__.co_argcount
        @wraps(f)
        def new_f(*args, **kwds):
            for (a, t) in zip(args, types):
                assert isinstance(a, t), \
                       "arg %r does not match %s" % (a,t)
            return f(*args, **kwds)
        return new_f
    return check_accepts

def returns(rtype):
    def check_returns(f):
        @wraps(f)
        def new_f(*args, **kwds):
            result = f(*args, **kwds)
            assert isinstance(result, rtype), \
                   "return value %r does not match %s" % (result,rtype)
            return result
        return new_f
    return check_returns

@accepts(int, (int,float))
@returns((int,float))
def func(arg1, arg2):
    return arg1 * arg2

I get an assertion error on assert len(types) == f.__code__.co_argcount.

A secondary question, if I may, is how could I extend the @accepts decorator to work on class methods, which may have the self or cls arguments.

Upvotes: 0

Views: 373

Answers (1)

LazyCoder
LazyCoder

Reputation: 1265

Change the decorations as below:

@returns((int,float))
@accepts(int, (int,float))
def func(arg1, arg2):
    return arg1 * arg2

Upvotes: 1

Related Questions