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?

Impi

Answer 1

Finally found an answer that works (with a minor change, see below). Full credits to @Kyuuhachi!

This assumes that we are using CPython. The key part is to use the _testcapi module.

The difference with @Kyuuhachi's answer is minor: in our case, we only want to remove the decorator itself from the stack trace. In other words, we want the same stack trace as if the function had not been decorated.

The decorator in my question becomes:

import functools
import sys
import _testcapi

def bar(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # ... (decorator functionality before calling func)
        try:
            result = func(*args, **kwargs)
        except:
            tp, exc, tb = sys.exc_info()
            _testcapi.set_exc_info(tp, exc, tb.tb_next)
            del tp, exc, tb
            raise

        # ... (decorator functionality after calling func)
        return result
    return wrapper

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

With that:

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

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

and

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

Input In [1], in f(x)
     21 @bar
     22 def f(x):
---> 23     return 1 / x

ZeroDivisionError: division by zero

Which is exactly what I wanted.

Impi