PRE и POST декораторы с аргументами? - PullRequest
Новая
       69

PRE и POST декораторы с аргументами?

1 голос
/ 12 марта 2020

Как мне написать обертку, которая добавляет и добавляет код на основе аргументов ... Вот оригинальный метод: 

 def method(self, arg1, arg2, arg3, ret='abc'):
     arg1 = pre_fun(arg1)
     rv = None
     ..... code ....
     if ret == 'abc' : return abc_fun(rv)
     if ret == 'efg' : return efg_fun(rv)

хочу преобразовать его во что-то вроде строки: 

  @pre(fun=pre_fun, arg='arg1')
  @post(ret1=abc_fun, arg_ret1='rv', ret2=efg_fun, arg_ret2='rv')
  def method(self, arg1, arg2, arg3, ret='abc'):
      rv = None
      ....... code .....

Я знаю, что это не совсем так. Является ли это возможным. Также я могу иметь значения по умолчанию, чтобы я мог сказать: 

  @pre
  @post
  def method(self, arg1, arg2, arg3, ret='abc'):
      rv = None
      ....... code .....

ИЛИ если нет Я в порядке в жестком кодировании параметров из get go. (Я бы даже предпочел это для краткости. Может быть даже @pre_post)

Я думаю, что мой arg_xxx = 'rv' немного ненадежен, но не могу понять иначе.

my работа в процессе, еще не проверена: 

def pp(fun):

   @functools.wraps(fun)
    def wrapper(*args, **kwargs):
        args[0] = xbitx(args[0])
        fun(*args, **kwargs)
        if ret == 'numpy' : return args[0]
        return iSDP(val=args[0], size=kwargs['size'], spaOnbits=kwargs['spaOnbits'])

    return wrapper

1 Ответ

1 голос
/ 20 марта 2020

Вы можете использовать фабрики декораторов и самоанализ, чтобы сделать что-то вроде того, что вы описали. Например, 

import inspect
import functools


def ensure_tuple(arg):
    if isinstance(arg, tuple):
        return arg
    else:
        return (arg,)

def default_pre_func(arg1):
    return 2*arg1

# the following function generates decorators
def pre(pre_func=default_pre_func):
    arg_spec = inspect.getfullargspec(pre_func)
    # build a decorator
    def pre_decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # you can do a lot more magic based on the arguments if you need/want to
            args = args
            num_args = len(arg_spec.args)  # the number of arguments that should be passed to pre_func
            new_args = ensure_tuple(pre_func(*args[:num_args])) + args[num_args:] 
            return func(*new_args, **kwargs)
        return wrapper
    # return this decorator
    return pre_decorator

def default_post_func1(rv):
    print("default_post_func1")
    return rv

def default_post_func2(rv):
    print("default_post_func2")
    return rv

def post(switch=None):
    """
    creates post_func decorators with the options defined in switch
    :param switch: a dictionary describing the actions
    returns: a decorator
    """
    if switch is None:
        switch = {
            1: default_post_func1,
            2: default_post_func2
        }
    # define the decorator
    def post_decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            rv, ret = func(*args, **kwargs)
            return switch[ret](rv)
        return wrapper
    return post_decorator

@pre()
@post()
def my_function(arg1, arg2):
    print(f"arg1: {arg1}")
    from random import choice
    ret = choice([1, 2])
    print(f"ret: {ret}")
    return arg2, ret

my_function(1,"a")

Редактировать: кстати, если вы хотите избавиться от скобок при использовании аргументов по умолчанию, есть хорошее описание того, как это сделать здесь .

...