def frames_to_traceback(
frame: Optional[FrameType],
limit: int,
should_prune: Optional[Callable[[str], bool]] = None,
) -> Optional[TracebackType]:
ctb: Optional[TracebackType] = None
skipped = False
while frame is not None and limit > 0:
if _MODIFIED_EXCEPTION_VAR_NAME in frame.f_locals:
return TracebackType(
tb_next=None,
tb_frame=frame,
tb_lasti=frame.f_lasti,
tb_lineno=frame.f_lineno,
)
if not skipped:
if should_prune is not None and should_prune(
frame.f_globals["__name__"]):
frame = frame.f_back
continue
skipped = True
if should_prune is None or not should_prune(
frame.f_globals["__name__"]):
ctb = TracebackType(
tb_next=ctb,
tb_frame=frame,
tb_lasti=frame.f_lasti,
tb_lineno=frame.f_lineno,
)
limit -= 1
frame = frame.f_back
continue
break
return ctb
def test_construct():
frame = sys._getframe()
tb = get_tb()
tb2 = TracebackType(tb, frame, 1, 2)
assert tb2.tb_next is tb
assert tb2.tb_frame is frame
assert tb2.tb_lasti == 1
assert tb2.tb_lineno == 2
tb2 = TracebackType(tb, frame, 1, -1)
assert tb2.tb_next is tb
assert tb2.tb_frame is frame
assert tb2.tb_lasti == 1
assert tb2.tb_lineno == -1
def modify_traceback(
traceback: Optional[TracebackType],
should_prune: Optional[Callable[[str], bool]] = None,
add_traceback: Optional[TracebackType] = None,
) -> Optional[TracebackType]:
ctb: Optional[TracebackType] = None
# get stack
stack: List[TracebackType] = []
if add_traceback is not None:
f: Optional[TracebackType] = add_traceback
while f is not None:
stack.append(f)
f = f.tb_next
f = traceback
while f is not None:
stack.append(f)
f = f.tb_next
stack.reverse()
# prune and reconstruct
for n, f in enumerate(stack):
if (n == 0 or should_prune is None
or not should_prune(f.tb_frame.f_globals["__name__"])):
ctb = TracebackType(
tb_next=ctb,
tb_frame=f.tb_frame,
tb_lasti=f.tb_lasti,
tb_lineno=f.tb_lineno,
)
return ctb
def equip_with_traceback(exc, depth=0): # Python 3.7+
"""Given an exception instance exc, equip it with a traceback.
`depth` is the starting depth for `sys._getframe`.
The return value is `exc`, with its traceback set to the produced
traceback.
Python 3.7 and later only.
When not supported, raises `NotImplementedError`.
This is useful in some special cases only, mainly when `raise` cannot be
used for some reason, and a manually created exception instance needs a
traceback. (E.g. in implementing the system for conditions and restarts.)
The `sys._getframe` function exists in CPython and in PyPy3,
but for another arbitrary Python implementation this is not
guaranteed.
Based on solution by StackOverflow user Zbyl:
https://stackoverflow.com/a/54653137
See also:
https://docs.python.org/3/library/types.html#types.TracebackType
https://docs.python.org/3/reference/datamodel.html#traceback-objects
https://docs.python.org/3/library/sys.html#sys._getframe
"""
if not isinstance(exc, BaseException):
raise TypeError(
"exc must be an exception instance; got {} with value '{}'".format(
type(exc), exc))
try:
getframe = sys._getframe
except AttributeError as err: # pragma: no cover, both CPython and PyPy3 have sys._getframe.
raise NotImplementedError(
"Need a Python interpreter which has `sys._getframe`") from err
tb = None
while True:
# Starting from given depth, get all frames up to the root of the call stack.
try:
frame = getframe(depth)
depth += 1
except ValueError:
break
# Python 3.7+ allows creating types.TracebackType objects from Python code.
try:
tb = TracebackType(tb, frame, frame.f_lasti, frame.f_lineno)
except TypeError as err: # Python 3.6 or earlier
raise NotImplementedError(
"Need Python 3.7 or later to create traceback objects") from err
return exc.with_traceback(tb) # Python 3.7+
def make_traceback():
tb = None
depth = 0
while True:
try:
frame = sys._getframe(depth)
depth += 1
except ValueError:
break
tb = TracebackType(tb, frame, frame.f_lasti, frame.f_lineno)
return tb
def name_error(varname, function, pop_frames=1):
"""Raise a PteraNameError pointing to the right location."""
fr = inspect.currentframe()
for i in range(pop_frames + 1):
if fr:
fr = fr.f_back
err = PteraNameError(varname, function)
try: # pragma: no cover
tb = TracebackType(
tb_next=None,
tb_frame=fr,
tb_lasti=fr.f_lasti,
tb_lineno=fr.f_lineno,
)
return err.with_traceback(tb)
except TypeError: # pragma: no cover
return err
def try_simplify_traceback(tb: TracebackType) -> Optional[TracebackType]:
"""
Simplify the traceback. It removes the tracebacks in the current package, and only
shows the traceback that is related to the thirdparty and user-specified codes.
Returns
-------
TracebackType or None
Simplified traceback instance. It returns None if it fails to simplify.
Notes
-----
This keeps the tracebacks once it sees they are from a different file even
though the following tracebacks are from the current package.
Examples
--------
>>> import importlib
>>> import sys
>>> import traceback
>>> import tempfile
>>> with tempfile.TemporaryDirectory() as tmp_dir:
... with open("%s/dummy_module.py" % tmp_dir, "w") as f:
... _ = f.write(
... 'def raise_stop_iteration():\\n'
... ' raise StopIteration()\\n\\n'
... 'def simple_wrapper(f):\\n'
... ' def wrapper(*a, **k):\\n'
... ' return f(*a, **k)\\n'
... ' return wrapper\\n')
... f.flush()
... spec = importlib.util.spec_from_file_location(
... "dummy_module", "%s/dummy_module.py" % tmp_dir)
... dummy_module = importlib.util.module_from_spec(spec)
... spec.loader.exec_module(dummy_module)
>>> def skip_doctest_traceback(tb):
... import pyspark
... root = os.path.dirname(pyspark.__file__)
... pairs = zip(walk_tb(tb), traceback.extract_tb(tb))
... for cur_tb, cur_frame in pairs:
... if cur_frame.filename.startswith(root):
... return cur_tb
Regular exceptions should show the file name of the current package as below.
>>> exc_info = None
>>> try:
... fail_on_stopiteration(dummy_module.raise_stop_iteration)()
... except Exception as e:
... tb = sys.exc_info()[-1]
... e.__cause__ = None
... exc_info = "".join(
... traceback.format_exception(type(e), e, tb))
>>> print(exc_info) # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
Traceback (most recent call last):
File ...
...
File "/.../pyspark/util.py", line ...
...
RuntimeError: ...
>>> "pyspark/util.py" in exc_info
True
If the traceback is simplified with this method, it hides the current package file name:
>>> exc_info = None
>>> try:
... fail_on_stopiteration(dummy_module.raise_stop_iteration)()
... except Exception as e:
... tb = try_simplify_traceback(sys.exc_info()[-1])
... e.__cause__ = None
... exc_info = "".join(
... traceback.format_exception(
... type(e), e, try_simplify_traceback(skip_doctest_traceback(tb))))
>>> print(exc_info) # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
RuntimeError: ...
>>> "pyspark/util.py" in exc_info
False
In the case below, the traceback contains the current package in the middle.
In this case, it just hides the top occurrence only.
>>> exc_info = None
>>> try:
... fail_on_stopiteration(dummy_module.simple_wrapper(
... fail_on_stopiteration(dummy_module.raise_stop_iteration)))()
... except Exception as e:
... tb = sys.exc_info()[-1]
... e.__cause__ = None
... exc_info_a = "".join(
... traceback.format_exception(type(e), e, tb))
... exc_info_b = "".join(
... traceback.format_exception(
... type(e), e, try_simplify_traceback(skip_doctest_traceback(tb))))
>>> exc_info_a.count("pyspark/util.py")
2
>>> exc_info_b.count("pyspark/util.py")
1
"""
if "pypy" in platform.python_implementation().lower():
# Traceback modification is not supported with PyPy in PySpark.
return None
if sys.version_info[:2] < (3, 7):
# Traceback creation is not supported Python < 3.7.
# See https://bugs.python.org/issue30579.
return None
import pyspark
root = os.path.dirname(pyspark.__file__)
tb_next = None
new_tb = None
pairs = zip(walk_tb(tb), traceback.extract_tb(tb))
last_seen = []
for cur_tb, cur_frame in pairs:
if not cur_frame.filename.startswith(root):
# Filter the stacktrace from the PySpark source itself.
last_seen = [(cur_tb, cur_frame)]
break
for cur_tb, cur_frame in reversed(list(itertools.chain(last_seen, pairs))):
# Once we have seen the file names outside, don't skip.
new_tb = TracebackType(
tb_next=tb_next,
tb_frame=cur_tb.tb_frame,
tb_lasti=cur_tb.tb_frame.f_lasti,
tb_lineno=cur_tb.tb_frame.f_lineno if cur_tb.tb_frame.f_lineno is not None else -1,
)
tb_next = new_tb
return new_tb
def tb_set_next(tb: TracebackType,
tb_next: t.Optional[TracebackType]) -> TracebackType:
tb.tb_next = tb_next
return tb
def equip_with_traceback(exc, stacklevel=1): # Python 3.7+
"""Given an exception instance exc, equip it with a traceback.
`stacklevel` is the starting depth below the top of the call stack,
to cull useless detail:
- `0` means the trace includes everything, also
`equip_with_traceback` itself,
- `1` means the trace includes everything up to the caller,
- And so on.
So typically, for direct use of this function `stacklevel` should
be `1` (so it excludes `equip_with_traceback` itself, but shows
all stack levels from your code), and for use in a utility function
that itself is called from your code, it should be `2` (so it excludes
the utility function, too).
The return value is `exc`, with its traceback set to the produced
traceback.
Python 3.7 and later only.
When not supported, raises `NotImplementedError`.
This is useful mainly in special cases, where `raise` cannot be used for
some reason, and a manually created exception instance needs a traceback.
(The `signal` function in the conditions-and-restarts system uses this.)
**CAUTION**: The `sys._getframe` function exists in CPython and in PyPy3,
but for another arbitrary Python implementation this is not guaranteed.
Based on solution by StackOverflow user Zbyl:
https://stackoverflow.com/a/54653137
See also:
https://docs.python.org/3/library/types.html#types.TracebackType
https://docs.python.org/3/reference/datamodel.html#traceback-objects
https://docs.python.org/3/library/sys.html#sys._getframe
"""
if not isinstance(exc, BaseException):
raise TypeError(f"exc must be an exception instance; got {type(exc)} with value {repr(exc)}")
if not isinstance(stacklevel, int):
raise TypeError(f"stacklevel must be int, got {type(stacklevel)} with value {repr(stacklevel)}")
if stacklevel < 0:
raise ValueError(f"stacklevel must be >= 0, got {repr(stacklevel)}")
try:
getframe = sys._getframe
except AttributeError as err: # pragma: no cover, both CPython and PyPy3 have sys._getframe.
raise NotImplementedError("Need a Python interpreter which has `sys._getframe`") from err
frames = []
depth = stacklevel
while True:
try:
frames.append(getframe(depth)) # 0 = top of call stack
depth += 1
except ValueError: # beyond the root level
break
# Python 3.7+ allows creating `types.TracebackType` objects in Python code.
try:
tracebacks = []
nxt = None # tb_next should point toward the level where the exception occurred.
for frame in frames: # walk from top of call stack toward the root
tb = TracebackType(nxt, frame, frame.f_lasti, frame.f_lineno)
tracebacks.append(tb)
nxt = tb
if tracebacks:
tb = tracebacks[-1] # root level
else:
tb = None
except TypeError as err: # Python 3.6 or earlier
raise NotImplementedError("Need Python 3.7 or later to create traceback objects") from err
return exc.with_traceback(tb) # Python 3.7+