How to use Python decorators to check function arg

问题:

I would like to define some generic decorators to check arguments before calling some functions.

Something like:

@checkArguments(types = ['int', 'float'])
def myFunction(thisVarIsAnInt, thisVarIsAFloat)
    ''' Here my code '''
    pass

Side notes:

1. Type checking is just here to show an example
2. I'm using Python 2.7 but Python 3.0 whould be interesting too

回答1:

From the Decorators for Functions and Methods:

def accepts(*types):
    def check_accepts(f):
        assert len(types) == f.func_code.co_argcount
        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)
        new_f.func_name = f.func_name
        return new_f
    return check_accepts

Usage:

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

func(3, 2) # -> 6
func('3', 2) # -> AssertionError: arg '3' does not match <type 'int'>

回答2:

On Python 3.3, you can use function annotations and inspect:

import inspect

def validate(f):
    def wrapper(*args):
        fname = f.__name__
        fsig = inspect.signature(f)
        vars = ', '.join('{}={}'.format(*pair) for pair in zip(fsig.parameters, args))
        params={k:v for k,v in zip(fsig.parameters, args)}
        print('wrapped call to {}({})'.format(fname, params))
        for k, v in fsig.parameters.items():
            p=params[k]
            msg='call to {}({}): {} failed {})'.format(fname, vars, k, v.annotation.__name__)
            assert v.annotation(params[k]), msg
        ret = f(*args)
        print('  returning {} with annotation: "{}"'.format(ret, fsig.return_annotation))
        return ret
    return wrapper

@validate
def xXy(x: lambda _x: 10<_x<100, y: lambda _y: isinstance(_y,float)) -> ('x times y','in X and Y units'):
    return x*y

xy = xXy(10,3)
print(xy)

If there is a validation error, prints:

AssertionError: call to xXy(x=12, y=3): y failed <lambda>)

If there is not a validation error, prints:

wrapped call to xXy({'y': 3.0, 'x': 12})
  returning 36.0 with annotation: "('x times y', 'in X and Y units')"

You can use a function rather than a lambda to get a name in the assertion failure.

回答3:

As you certainly know, it's not pythonic to reject an argument only based on its type.
Pythonic approach is rather "try to deal with it first"
That's why I would rather do a decorator to convert the arguments

def enforce(*types):
    def decorator(f):
        def new_f(*args, **kwds):
            #we need to convert args into something mutable   
            newargs = []        
            for (a, t) in zip(args, types):
               newargs.append( t(a)) #feel free to have more elaborated convertion
            return f(*newargs, **kwds)
        return new_f
    return decorator

This way, your function is fed with the type you expect But if the parameter can quack like a float, it is accepted

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

print (func(3, 2)) # -> 6.0
print (func('3', 2)) # -> 6.0
print (func('three', 2)) # -> ValueError: invalid literal for int() with base 10: 'three'

I use this trick (with proper conversion method) to deal with vectors.
Many methods I write expect MyVector class as it has plenty of functionalities; but sometime you just want to write

transpose ((2,4))

回答4:

To enforce string arguments to a parser that would throw cryptic errors when provided with non-string input, I wrote the following, which tries to avoid allocation and function calls:

from functools import wraps

def argtype(**decls):
    """Decorator to check argument types.

    Usage:

    @argtype(name=str, text=str)
    def parse_rule(name, text): ...
    """

    def decorator(func):
        code = func.func_code
        fname = func.func_name
        names = code.co_varnames[:code.co_argcount]

        @wraps(func)
        def decorated(*args,**kwargs):
            for argname, argtype in decls.iteritems():
                try:
                    argval = args[names.index(argname)]
                except ValueError:
                    argval = kwargs.get(argname)
                if argval is None:
                    raise TypeError("%s(...): arg '%s' is null"
                                    % (fname, argname))
                if not isinstance(argval, argtype):
                    raise TypeError("%s(...): arg '%s': type is %s, must be %s"
                                    % (fname, argname, type(argval), argtype))
            return func(*args,**kwargs)
        return decorated

    return decorator

回答5:

All of these posts seem out of date - pint now provides this functionality built in. See here. Copied here for posterity:

Checking dimensionality When you want pint quantities to be used as inputs to your functions, pint provides a wrapper to ensure units are of correct type - or more precisely, they match the expected dimensionality of the physical quantity.

Similar to wraps(), you can pass None to skip checking of some parameters, but the return parameter type is not checked.

>>> mypp = ureg.check('[length]')(pendulum_period) 

In the decorator format:

>>> @ureg.check('[length]')
... def pendulum_period(length):
...     return 2*math.pi*math.sqrt(length/G)

回答6:

I have a slightly improved version of @jbouwmans sollution, using python decorator module, which makes the decorator fully transparent and keeps not only signature but also docstrings in place and might be the most elegant way of using decorators

from decorator import decorator

def check_args(**decls):
    """Decorator to check argument types.

    Usage:

    @check_args(name=str, text=str)
    def parse_rule(name, text): ...
    """
    @decorator
    def wrapper(func, *args, **kwargs):
        code = func.func_code
        fname = func.func_name
        names = code.co_varnames[:code.co_argcount]
        for argname, argtype in decls.iteritems():
            try:
                argval = args[names.index(argname)]
            except IndexError:
                argval = kwargs.get(argname)
            if argval is None:
                raise TypeError("%s(...): arg '%s' is null"
                            % (fname, argname))
            if not isinstance(argval, argtype):
                raise TypeError("%s(...): arg '%s': type is %s, must be %s"
                            % (fname, argname, type(argval), argtype))
    return func(*args, **kwargs)
return wrapper

回答7:

I think the Python 3.5 answer to this question is beartype. As explained in this post it comes with handy features. Your code would then look like this

from beartype import beartype
@beartype
def sprint(s: str) -> None:
   print(s)

and results in

>>> sprint("s")
s
>>> sprint(3)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<string>", line 13, in func_beartyped
TypeError: sprint() parameter s=3 not of <class 'str'>