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:
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?
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.