How to "decorate" a function in the call site rather than definition

Question

I work on the useful introspection helpers on the exception, thus wanted to generalize the code rather than copy-pasting it all the time. I have

def state_on_exc(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        try:
            result = f(*args, **kwargs)
            return result
        except Exception as e:
            ex_type, exc_value, tb = sys.exc_info()
            if tb is not None:
                prev = tb
                curr = tb.tb_next
                while curr is not None:
                    prev = curr
                    curr = curr.tb_next
                print(prev.tb_frame.f_locals)
            raise e

@state_on_exc
def apply(f):
    return f

def myfunc():
    a = 5
    raise ValueError

apply(myfunc())

But the wrapper doesn't seem to be called on exception (it raised specifically in the myfunc()) - it doesn't print any local variables. Is there any proper way to achieve the same, or any better way to do that?

Answer 1

Quite close.

1. You need to return the wrapped function from the decorator.
2. To re-raise, you should not add the exception object to the raise stmt.
3. To decorate the function, pass the function object to the decorator function.
4. To call the decorated function, call the decorated function. And not the non-decorated function.

Example:

from functools import wraps

def decorate(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except Exception as e:
            # handle exception
            print("handled exception:", type(e))
            # re-raise
            raise
    return wrapper

def myfunc():
    a = 5
    raise ValueError

decorated_fun = decorate(myfunc)
decorated_fun()

The output is:

$ python3 test.py
handled exception: <class 'ValueError'>
Traceback (most recent call last):
  File "tt.py", line 25, in <module>
    decorated_fun()
  File "tt.py", line 7, in wrapper
    result = f(*args, **kwargs)
  File "tt.py", line 22, in myfunc
    raise ValueError
ValueError