Как передать определенный аргумент декоратору в python - PullRequest
3 голосов
/ 30 июня 2010

Я хочу написать декоратор функции Python, который проверяет, соответствуют ли определенные аргументы функции некоторому критерию.Например, предположим, что я хочу проверить, что некоторые аргументы всегда равны, тогда я хочу иметь возможность сделать что-то вроде этого (не допустимый код Python)

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

Но я не могу понять, как добиться выше,Есть ли способ передать конкретные аргументы декоратору?(например, arg2 для foo и arg3 для bar выше)

Спасибо!

Ответы [ 2 ]

11 голосов
/ 30 июня 2010

Вы можете сделать это:

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

Итак:

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

test(1,2) #succeeds
test(1,3) #fails
4 голосов
/ 01 июля 2010

Мой предыдущий ответ пропустил суть вашего вопроса.

Вот более длинное решение:

def ensure_even(*argvars):
  def fdec(func):
    def f(*args,**kwargs):
      for argvar in argvars:
        try:
          assert(not args[func.func_code.co_varnames.index(argvar)] % 2)
        except IndexError:
          assert(not kwargs[argvar] % 2)
      return func(*args,**kwargs)
    return f
  return fdec

, и поэтому это использование:

@ensure_even('a','b')
def both_even(a,b):
    print a,b

@ensure_even('even')
def first_even(even, arg2):
    print even, arg2

both_even(2,2)
first_even(2,3)
both_even(2,1) #fails

хотяЯ не уверен, что это будет работать во всех ситуациях.

...