带参数的 PRE 和 POST 装饰器？

Question

如何编写一个基于参数添加和添加代码的包装器。 这是原始方法：

 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 .....

或者如果不是我可以硬编码从一开始就设置参数。 （我什至更喜欢它的简短性。甚至可能是@pre_post）

我认为我的 arg_xxx='rv' 有点不稳定，但无法想象其他方式。

---

我的工作正在进行中，尚未测试：

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  

Answer 1

您可以使用装饰器工厂和自省来执行您所描述的操作。例如

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")

编辑：顺便说一句，如果你想在使用默认参数时去掉括号，有一个很好的描述如何做到这一点here。