Simplify removing internal tracebacks from exceptions with cause. (#460)

docs/source/changes.md

@@ -15,6 +15,7 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 - {pull}`456` refers to the source code on Github when clicking on a source link.
 - {pull}`457` refactors everything around formatting node names.
 - {pull}`459` adds a pre-commit hook to sort `__all__`.
+- {pull}`460` simplifies removing internal tracebacks from exceptions with a cause.
 
 ## 0.4.1 - 2023-10-11
 

src/_pytask/dag.py

@@ -33,7 +33,6 @@
 from _pytask.nodes import PythonNode
 from _pytask.report import DagReport
 from _pytask.shared import reduce_names_of_multiple_nodes
-from _pytask.traceback import remove_internal_traceback_frames_from_exception
 from _pytask.traceback import render_exc_info
 from _pytask.tree_util import tree_map
 from rich.text import Text
@@ -242,7 +241,6 @@ def _check_if_root_nodes_are_available(dag: nx.DiGraph, paths: Sequence[Path]) -
                 try:
                     node_exists = dag.nodes[node]["node"].state()
                 except Exception as e:  # noqa: BLE001
-                    e = remove_internal_traceback_frames_from_exception(e)
                     msg = _format_exception_from_failed_node_state(node, dag)
                     raise ResolvingDependenciesError(msg) from e
                 if not node_exists:

src/_pytask/execute.py

@@ -33,7 +33,6 @@
 from _pytask.outcomes import WouldBeExecuted
 from _pytask.report import ExecutionReport
 from _pytask.traceback import format_exception_without_traceback
-from _pytask.traceback import remove_internal_traceback_frames_from_exception
 from _pytask.traceback import remove_traceback_from_exc_info
 from _pytask.traceback import render_exc_info
 from _pytask.tree_util import tree_leaves
@@ -149,7 +148,6 @@ def _safe_load(node: PNode, task: PTask) -> Any:
     try:
         return node.load()
     except Exception as e:  # noqa: BLE001
-        e = remove_internal_traceback_frames_from_exception(e)
         task_name = getattr(task, "display_name", task.name)
         msg = f"Exception while loading node {node.name!r} of task {task_name!r}"
         raise NodeLoadError(msg) from e

src/_pytask/traceback.py

@@ -21,7 +21,6 @@
 __all__ = [
     "format_exception_without_traceback",
     "remove_internal_traceback_frames_from_exc_info",
-    "remove_internal_traceback_frames_from_exception",
     "remove_traceback_from_exc_info",
     "render_exc_info",
 ]
@@ -65,13 +64,18 @@ def remove_traceback_from_exc_info(
     return (exc_info[0], exc_info[1], None)  # type: ignore[return-value]
 
 
-def remove_internal_traceback_frames_from_exception(exc: Exception) -> Exception:
+def _remove_internal_traceback_frames_from_exception(
+    exc: BaseException | None,
+) -> BaseException | None:
     """Remove internal traceback frames from exception.
 
     The conversion between exceptions and ``sys.exc_info`` is explained here:
     https://stackoverflow.com/a/59041463/7523785.
 
     """
+    if exc is None:
+        return exc
+
     _, _, tb = remove_internal_traceback_frames_from_exc_info(
         (type(exc), exc, exc.__traceback__)
     )
@@ -88,6 +92,11 @@ def remove_internal_traceback_frames_from_exc_info(
     occurrence downwards.
 
     """
+    if isinstance(exc_info[1], Exception):
+        exc_info[1].__cause__ = _remove_internal_traceback_frames_from_exception(
+            exc_info[1].__cause__
+        )
+
     if isinstance(exc_info[2], TracebackType):
         filtered_traceback = _filter_internal_traceback_frames(exc_info)
         exc_info = (*exc_info[:2], filtered_traceback)

tests/test_dag.py

@@ -213,6 +213,10 @@ def task_example(a = PythonNode(value={"a": 1}, hash=True)):
     assert result.exit_code == ExitCode.DAG_FAILED
     assert "task_example" in result.output
 
+    # Assert that the traceback is hidden.
+    assert "_pytask/nodes.py" not in result.output
+    assert "in state" not in result.output
+
 
 def test_python_nodes_are_unique(tmp_path):
     tmp_path.joinpath("a").mkdir()

tests/test_execute.py

@@ -832,4 +832,6 @@ def task_example(
     assert result.exit_code == ExitCode.FAILED
     assert "task_example.py::task_example" in result.output
     assert "Exception while loading node" in result.output
+
+    # Test that traceback is hidden.
     assert "_pytask/execute.py" not in result.output