How to use Python decorators to check function arg

2019-01-06 12:04发布

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

7条回答
虎瘦雄心在
2楼-- · 2019-01-06 12:27

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
查看更多
混吃等死
3楼-- · 2019-01-06 12:30

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'>
查看更多
Lonely孤独者°
4楼-- · 2019-01-06 12:33

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'>
查看更多
SAY GOODBYE
5楼-- · 2019-01-06 12:34

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楼-- · 2019-01-06 12:36

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
查看更多
劫难
7楼-- · 2019-01-06 12:46

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.

查看更多
登录 后发表回答