Question

I want to write a python function decorator that tests that certain arguments to a function pass some criterion. For eg, Suppose I want to test that some arguments are always even, then I want to be able to do something like this (not valid python code)

def ensure_even( n )  :
  def decorator( function ) :
    @functools.wraps( function )
    def wrapper(*args, **kwargs):
      assert(n % 2 == 0)
      return function(*args, **kwargs)
    return wrapper
   return decorator


@ensure_even(arg2)
def foo(arg1, arg2, arg3) : pass

@ensure_even(arg3)
def bar(arg1, arg2, arg3) : pass

But I cannot figure how to achieve the above. Is there a way to pass specific arguments to the decorator? (like arg2 for foo and arg3 for bar in the above)

Thanks!

Answer 1

You can do this:

def ensure_even(argnum):
  def fdec(func):
    def f(*args, **kwargs):
      assert(args[argnum] % 2 == 0)  #or assert(not args[argnum] % 2)
      return func(*args, **kwargs)
    return f
  return fdec

So then:

@ensure_even(1)  #2nd argument must be even
def test(arg1, arg2):
  print(arg2)

test(1,2) #succeeds
test(1,3) #fails