带参数的PRE和POST装饰器?

时间:2020-03-12 00:55:54

标签: python post decorator wrapper pre

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

 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  

1 个答案:

答案 0 :(得分: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