From 72f7289fc8a0d17f23dca0f9548fbc5d0962adbc Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Tue, 19 Mar 2024 20:20:16 +0100
Subject: [PATCH] Backport #580 to v0.4.7. (#580)

---
 docs/source/changes.md   |  4 ++++
 src/_pytask/debugging.py |  6 +++--
 tests/test_debugging.py  | 47 ++++++++++++++++++++++++++++++++++++----
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/docs/source/changes.md b/docs/source/changes.md
index 3a0a8270..3528db2f 100644
--- a/docs/source/changes.md
+++ b/docs/source/changes.md
@@ -5,6 +5,10 @@ chronological order. Releases follow [semantic versioning](https://semver.org/)
 releases are available on [PyPI](https://pypi.org/project/pytask) and
 [Anaconda.org](https://anaconda.org/conda-forge/pytask).
 
+## 0.4.7 - 2024-03-19
+
+- {pull}`580` is a backport of {pull}`579`.
+
 ## 0.4.6 - 2024-03-13
 
 - {pull}`576` fixes accidentally collecting `pytask.MarkGenerator` when using
diff --git a/src/_pytask/debugging.py b/src/_pytask/debugging.py
index c95f79d3..5d5c9e1b 100644
--- a/src/_pytask/debugging.py
+++ b/src/_pytask/debugging.py
@@ -328,7 +328,7 @@ def wrapper(*args: Any, **kwargs: Any) -> None:
         capman = session.config["pm"].get_plugin("capturemanager")
         live_manager = session.config["pm"].get_plugin("live_manager")
         try:
-            task_function(*args, **kwargs)
+            return task_function(*args, **kwargs)
 
         except Exception:
             # Order is important! Pausing the live object before the capturemanager
@@ -409,11 +409,13 @@ def wrapper(*args: Any, **kwargs: Any) -> None:
             console.rule("Captured stderr", style="default")
             console.print(err)
 
-        _pdb.runcall(task_function, *args, **kwargs)
+        out = _pdb.runcall(task_function, *args, **kwargs)
 
         live_manager.resume()
         capman.resume()
 
+        return out
+
     task.function = wrapper
 
 
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 77802a62..3dd49599 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -21,6 +21,12 @@
     IS_PEXPECT_INSTALLED = True
 
 
+def _escape_ansi(line):
+    """Escape ANSI sequences produced by rich."""
+    ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
+    return ansi_escape.sub("", line)
+
+
 @pytest.mark.unit()
 @pytest.mark.parametrize(
     ("value", "expected", "expectation"),
@@ -487,7 +493,40 @@ def test_function():
     _flush(child)
 
 
-def _escape_ansi(line):
-    """Escape ANSI sequences produced by rich."""
-    ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
-    return ansi_escape.sub("", line)
+@pytest.mark.end_to_end()
+@pytest.mark.skipif(not IS_PEXPECT_INSTALLED, reason="pexpect is not installed.")
+@pytest.mark.skipif(sys.platform == "win32", reason="pexpect cannot spawn on Windows.")
+def test_pdb_with_task_that_returns(tmp_path, runner):
+    source = """
+    from typing_extensions import Annotated
+    from pathlib import Path
+
+    def task_example() -> Annotated[str, Path("data.txt")]:
+        return "1"
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+
+    result = runner.invoke(cli, [tmp_path.as_posix(), "--pdb"])
+    assert result.exit_code == ExitCode.OK
+    assert tmp_path.joinpath("data.txt").read_text() == "1"
+
+
+@pytest.mark.end_to_end()
+@pytest.mark.skipif(not IS_PEXPECT_INSTALLED, reason="pexpect is not installed.")
+@pytest.mark.skipif(sys.platform == "win32", reason="pexpect cannot spawn on Windows.")
+def test_trace_with_task_that_returns(tmp_path):
+    source = """
+    from typing_extensions import Annotated
+    from pathlib import Path
+
+    def task_example() -> Annotated[str, Path("data.txt")]:
+        return "1"
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+
+    child = pexpect.spawn(f"pytask {tmp_path.as_posix()}")
+    child.sendline("c")
+    rest = child.read().decode("utf8")
+    assert "1  Succeeded" in _escape_ansi(rest)
+    assert tmp_path.joinpath("data.txt").read_text() == "1"
+    _flush(child)