PRE and POST decorators with arguments?

Question

How do I write a wrapper that prepend and append code based on arguments.. Here is the original method :

 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)

want to convert it to something along the line :

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

I know it is not exactly like this. Is it possible. Also can I have defaults, so that I can say :

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

OR if not I'm OK in hard-coding the parameters from the get go. (I would even prefer it for shortness.May be even @pre_post )

I think my arg_xxx='rv' is a bit flaky , but cant figure other way.

my work in progress, not tested yet :

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

Does https://stackoverflow.com/questions/739654/how-to-make-a-chain-of-function-decorators, specifically https://stackoverflow.com/a/1594484 help? — Karl Knechtel, Mar 12 '20 at 01:03

simon · Answer 1 · 2020-03-19T23:43:30.373

You could use decorator factories and introspection to do something like what you described. E.g.

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

Edit: btw if you want to get rid of the brackets when using default arguments, there's a nice description of how to do that here.

PRE and POST decorators with arguments?

1 Answers1