remove decorator from stack trace

Question

I want to write a function decorator that removes itself from the stack trace when an exception occurs (outside the logic specific to the decorator itself), for example when:

• the caller uses arguments that don't match the function signature, or
• the decorated function itself raises an exception.

Consider the following example:

import functools

def foo(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # ... (decorator functionality before calling func)
        result = func(*args, **kwargs)
        # ... (decorator functionality after calling func)
        return result
    return wrapper

@foo
def f(x):
    return 1 / x

Unfortunately:

>>> f()
TypeError                                 Traceback (most recent call last)
Input In [2], in <cell line: 1>()
----> 1 f()

Input In [1], in foo.<locals>.wrapper(*args, **kwargs)
      4 @functools.wraps(func)
      5 def wrapper(*args, **kwargs):
      6     # ... (decorator functionality before calling func)
----> 7     result = func(*args, **kwargs)
      8     # ... (decorator functionality after calling func)
      9     return result

TypeError: f() missing 1 required positional argument: 'x'

Likewise:

>>> f(0)
ZeroDivisionError                         Traceback (most recent call last)
Input In [3], in <cell line: 1>()
----> 1 f(0)

Input In [1], in foo.<locals>.wrapper(*args, **kwargs)
      4 @functools.wraps(func)
      5 def wrapper(*args, **kwargs):
      6     # ... (decorator functionality before calling func)
----> 7     result = func(*args, **kwargs)
      8     # ... (decorator functionality after calling func)
      9     return result

Input In [1], in f(x)
     12 @foo
     13 def f(x):
---> 14     return 1 / x

ZeroDivisionError: division by zero

This leads to "polluted" stack traces that include the decorator code context, file, lineno etc. The problem is compounded when we have nested decorated functions.

By contrast, observe how e.g. lru_cache keeps the traceback clean:

@functools.lru_cache(maxsize=4)
def f(x):
    return 1 / x

>>> f()
TypeError                                 Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 f()

TypeError: f() missing 1 required positional argument: 'x'

>>> f(0)
ZeroDivisionError                         Traceback (most recent call last)
Input In [6], in <cell line: 1>()
----> 1 f(0)

Input In [4], in f(x)
      1 @functools.lru_cache(maxsize=4)
      2 def f(x):
----> 3     return 1 / x

ZeroDivisionError: division by zero

How to achieve similar cleanliness in custom decorators?

Answer 1

The trick that functools.lru_cache uses is to retrieve the unwrapped function using the .__wrapped__ property of the wrapped function as demonstrated here.

You can look at the CPython implementation of lru_cache where the specific technique of using .__wrapped__ done in the called update_wrapper.