From c6ff1156f8a3ec4dabc9634010a69244a2e22363 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Tue, 19 Dec 2023 23:20:31 +0100
Subject: [PATCH 1/6] Support partialed functions.

---
 .github/workflows/main.yml     |  4 +--
 .gitignore                     |  1 +
 CHANGES.md                     |  1 +
 src/pytask_parallel/execute.py | 64 +++++++++++++++++++---------------
 tests/test_execute.py          | 40 +++++++++++++++++----
 tox.ini                        |  2 +-
 6 files changed, 74 insertions(+), 38 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 053d9e2..223d7d7 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -42,7 +42,7 @@ jobs:
 
       - name: Run unit tests and doctests.
         shell: bash -l {0}
-        run: tox -e test -- tests -m "unit or (not integration and not end_to_end)" --cov=./ --cov-report=xml
+        run: tox -e test -- tests -m "unit or (not integration and not end_to_end)" --cov=src --cov=tests --cov-report=xml
 
       - name: Upload coverage report for unit tests and doctests.
         if: runner.os == 'Linux' && matrix.python-version == '3.10'
@@ -51,7 +51,7 @@ jobs:
 
       - name: Run end-to-end tests.
         shell: bash -l {0}
-        run: tox -e test -- tests -m end_to_end --cov=./ --cov-report=xml
+        run: tox -e test -- tests -m end_to_end --cov=src --cov=tests --cov-report=xml
 
       - name: Upload coverage reports of end-to-end tests.
         if: runner.os == 'Linux' && matrix.python-version == '3.10'
diff --git a/.gitignore b/.gitignore
index fd09476..7042fe6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,4 @@ __pycache__
 build
 dist
 src/pytask_parallel/_version.py
+tests/test_jupyter/*.txt
diff --git a/CHANGES.md b/CHANGES.md
index 1210989..3f4ed66 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -10,6 +10,7 @@ releases are available on [PyPI](https://pypi.org/project/pytask-parallel) and
 - {pull}`72` moves the project to `pyproject.toml`.
 - {pull}`75` updates the release strategy.
 - {pull}`79` add tests for Jupyter and fix parallelization with `PythonNode`s.
+- {pull}`80` adds support for partialed functions.
 
 ## 0.4.0 - 2023-10-07
 
diff --git a/src/pytask_parallel/execute.py b/src/pytask_parallel/execute.py
index c063343..ce6d642 100644
--- a/src/pytask_parallel/execute.py
+++ b/src/pytask_parallel/execute.py
@@ -6,6 +6,7 @@
 import time
 import warnings
 from concurrent.futures import Future
+from functools import partial
 from pathlib import Path
 from types import ModuleType
 from types import TracebackType
@@ -296,23 +297,7 @@ def _execute_task(  # noqa: PLR0913
                 exc_info, show_locals, console_options
             )
         else:
-            if "return" in task.produces:
-                structure_out = tree_structure(out)
-                structure_return = tree_structure(task.produces["return"])
-                # strict must be false when none is leaf.
-                if not structure_return.is_prefix(structure_out, strict=False):
-                    msg = (
-                        "The structure of the return annotation is not a subtree of "
-                        "the structure of the function return.\n\nFunction return: "
-                        f"{structure_out}\n\nReturn annotation: {structure_return}"
-                    )
-                    raise ValueError(msg)
-
-                nodes = tree_leaves(task.produces["return"])
-                values = structure_return.flatten_up_to(out)
-                for node, value in zip(nodes, values):
-                    node.save(value)
-
+            _handle_task_function_return(task, out)
             processed_exc_info = None
 
         task_display_name = getattr(task, "display_name", task.name)
@@ -347,6 +332,27 @@ def _process_exception(
     return (*exc_info[:2], text)
 
 
+def _handle_task_function_return(task: PTask, out: Any) -> None:
+    if "return" not in task.produces:
+        return
+
+    structure_out = tree_structure(out)
+    structure_return = tree_structure(task.produces["return"])
+    # strict must be false when none is leaf.
+    if not structure_return.is_prefix(structure_out, strict=False):
+        msg = (
+            "The structure of the return annotation is not a subtree of "
+            "the structure of the function return.\n\nFunction return: "
+            f"{structure_out}\n\nReturn annotation: {structure_return}"
+        )
+        raise ValueError(msg)
+
+    nodes = tree_leaves(task.produces["return"])
+    values = structure_return.flatten_up_to(out)
+    for node, value in zip(nodes, values):
+        node.save(value)
+
+
 class DefaultBackendNameSpace:
     """The name space for hooks related to threads."""
 
@@ -362,13 +368,13 @@ def pytask_execute_task(session: Session, task: Task) -> Future[Any] | None:
         if session.config["n_workers"] > 1:
             kwargs = _create_kwargs_for_task(task)
             return session.config["_parallel_executor"].submit(
-                _mock_processes_for_threads, func=task.execute, **kwargs
+                _mock_processes_for_threads, task=task, **kwargs
             )
         return None
 
 
 def _mock_processes_for_threads(
-    func: Callable[..., Any], **kwargs: Any
+    task: PTask, **kwargs: Any
 ) -> tuple[
     None, list[Any], tuple[type[BaseException], BaseException, TracebackType] | None
 ]:
@@ -381,10 +387,11 @@ def _mock_processes_for_threads(
     """
     __tracebackhide__ = True
     try:
-        func(**kwargs)
+        out = task.function(**kwargs)
     except Exception:  # noqa: BLE001
         exc_info = sys.exc_info()
     else:
+        _handle_task_function_return(task, out)
         exc_info = None
     return None, [], exc_info
 
@@ -430,18 +437,17 @@ def sleep(self) -> None:
 def _get_module(func: Callable[..., Any], path: Path | None) -> ModuleType:
     """Get the module of a python function.
 
-    For Python <3.10, functools.partial does not set a `__module__` attribute which is
-    why ``inspect.getmodule`` returns ``None`` and ``cloudpickle.pickle_by_value``
-    fails. In later versions, ``functools`` is returned and everything seems to work
-    fine.
+    ``functools.partial`` obfuscates the module of the function and
+    ``inspect.getmodule`` returns :mod`functools`. Therefore, we recover the original
+    function.
 
-    Therefore, we use the path from the task module to aid the search which works for
-    Python <3.10.
-
-    We do not unwrap the partialed function with ``func.func``, since pytask in general
-    does not really support ``functools.partial``. Instead, use ``@task(kwargs=...)``.
+    We use the path from the task module to aid the search although it is not clear
+    whether it helps.
 
     """
+    if isinstance(func, partial):
+        func = func.func
+
     if path:
         return inspect.getmodule(func, path.as_posix())
     return inspect.getmodule(func)
diff --git a/tests/test_execute.py b/tests/test_execute.py
index 49ec3d1..9786c67 100644
--- a/tests/test_execute.py
+++ b/tests/test_execute.py
@@ -233,10 +233,12 @@ def task_example() -> Annotated[str, Path("file.txt")]:
     """
     tmp_path.joinpath("task_example.py").write_text(textwrap.dedent(source))
     result = runner.invoke(
-        cli, [tmp_path.as_posix(), "--parallel-backend", parallel_backend]
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
     )
     assert result.exit_code == ExitCode.OK
-    assert tmp_path.joinpath("file.txt").exists()
+    assert (
+        tmp_path.joinpath("file.txt").read_text() == "Hello, Darkness, my old friend."
+    )
 
 
 @pytest.mark.end_to_end()
@@ -252,10 +254,12 @@ def test_task_without_path_that_return(runner, tmp_path, parallel_backend):
     """
     tmp_path.joinpath("task_example.py").write_text(textwrap.dedent(source))
     result = runner.invoke(
-        cli, [tmp_path.as_posix(), "--parallel-backend", parallel_backend]
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
     )
     assert result.exit_code == ExitCode.OK
-    assert tmp_path.joinpath("file.txt").exists()
+    assert (
+        tmp_path.joinpath("file.txt").read_text() == "Hello, Darkness, my old friend."
+    )
 
 
 @pytest.mark.end_to_end()
@@ -264,7 +268,8 @@ def test_task_without_path_that_return(runner, tmp_path, parallel_backend):
 def test_parallel_execution_is_deactivated(runner, tmp_path, flag, parallel_backend):
     tmp_path.joinpath("task_example.py").write_text("def task_example(): pass")
     result = runner.invoke(
-        cli, [tmp_path.as_posix(), "-n 2", "--parallel-backend", parallel_backend, flag]
+        cli,
+        [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend, flag],
     )
     assert result.exit_code == ExitCode.OK
     assert "Started 2 workers" not in result.output
@@ -278,7 +283,30 @@ def test_parallel_execution_is_deactivated(runner, tmp_path, flag, parallel_back
 def test_raise_error_on_breakpoint(runner, tmp_path, code, parallel_backend):
     tmp_path.joinpath("task_example.py").write_text(f"def task_example(): {code}")
     result = runner.invoke(
-        cli, [tmp_path.as_posix(), "-n 2", "--parallel-backend", parallel_backend]
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
     )
     assert result.exit_code == ExitCode.FAILED
     assert "You cannot use 'breakpoint()'" in result.output
+
+
+@pytest.mark.end_to_end()
+@pytest.mark.parametrize("parallel_backend", PARALLEL_BACKENDS)
+def test_task_partialed(runner, tmp_path, parallel_backend):
+    source = """
+    from pathlib import Path
+    from pytask import task
+    from functools import partial
+
+    def create_text(text):
+        return text
+
+    task_example = task(
+        produces=Path("file.txt")
+    )(partial(create_text, text="Hello, Darkness, my old friend."))
+    """
+    tmp_path.joinpath("task_example.py").write_text(textwrap.dedent(source))
+    result = runner.invoke(
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
+    )
+    assert result.exit_code == ExitCode.OK
+    assert tmp_path.joinpath("file.txt").exists()
diff --git a/tox.ini b/tox.ini
index 76e5c5e..c23cf3c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -3,7 +3,7 @@ requires = tox>=4
 envlist = test
 
 [testenv]
-package = wheel
+package = editable
 
 [testenv:test]
 extras = test

From a0d6c37b4d5b024e18bf2a98694a38a22a9875c0 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Sat, 30 Dec 2023 12:52:12 +0100
Subject: [PATCH 2/6] Fix testing for pytask v0.4.5.

---
 .gitignore        | 1 +
 pyproject.toml    | 4 ++--
 tests/conftest.py | 2 ++
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index fd09476..4a85b00 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,4 @@ __pycache__
 build
 dist
 src/pytask_parallel/_version.py
+tests/test_jupyter/file.txt
diff --git a/pyproject.toml b/pyproject.toml
index d9a0995..5910540 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,8 +52,8 @@ Tracker = "https://github.com/pytask-dev/pytask-parallel/issues"
 [tool.setuptools]
 include-package-data = true
 zip-safe = false
-platforms = [ "any",]
-license-files = [ "LICENSE",]
+platforms = ["any"]
+license-files = ["LICENSE"]
 
 [tool.check-manifest]
 ignore = ["src/pytask_parallel/_version.py"]
diff --git a/tests/conftest.py b/tests/conftest.py
index 8cba383..3881664 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -6,6 +6,7 @@
 
 import pytest
 from click.testing import CliRunner
+from pytask import storage
 
 
 class SysPathsSnapshot:
@@ -62,6 +63,7 @@ def _restore_sys_path_and_module_after_test_execution():
 class CustomCliRunner(CliRunner):
     def invoke(self, *args, **kwargs):
         """Restore sys.path and sys.modules after an invocation."""
+        storage.create()
         with restore_sys_path_and_module_after_test_execution():
             return super().invoke(*args, **kwargs)
 

From 37ed8ff163a1412a05323a8d265c4daa3ec0d2ae Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Sat, 30 Dec 2023 15:41:48 +0100
Subject: [PATCH 3/6] Disable errors with nbmake on macos.

---
 tests/conftest.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/conftest.py b/tests/conftest.py
index 3881664..ef47c25 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import os
 import sys
 from contextlib import contextmanager
 from typing import Callable
 
 import pytest
 from click.testing import CliRunner
+from nbmake.pytest_items import NotebookItem
 from pytask import storage
 
 
@@ -71,3 +73,11 @@ def invoke(self, *args, **kwargs):
 @pytest.fixture()
 def runner():
     return CustomCliRunner()
+
+
+def pytest_collection_modifyitems(session, config, items) -> None:  # noqa: ARG001
+    """Add markers to Jupyter notebook tests."""
+    if sys.platform == "darwin" and "CI" in os.environ:  # pragma: no cover
+        for item in items:
+            if isinstance(item, NotebookItem):
+                item.add_marker(pytest.mark.xfail(reason="Fails regularly on MacOS"))

From c53b7d434f75ee9e0aba06c25374548ea7ca1896 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Fri, 12 Jan 2024 16:09:53 +0100
Subject: [PATCH 4/6] Fix new pytask version.

---
 .pre-commit-config.yaml | 2 +-
 environment.yml         | 2 +-
 pyproject.toml          | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5857bd5..347348a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -65,7 +65,7 @@ repos:
             attrs,
             cloudpickle,
             loky,
-            pytask>=0.4.0,
+            pytask>=0.4.5,
             rich,
             types-click,
             types-setuptools,
diff --git a/environment.yml b/environment.yml
index 7b2ab71..27498f2 100644
--- a/environment.yml
+++ b/environment.yml
@@ -12,7 +12,7 @@ dependencies:
   - toml
 
   # Package dependencies
-  - pytask>=0.4.0
+  - pytask>=0.4.5
   - cloudpickle
   - loky
   - optree
diff --git a/pyproject.toml b/pyproject.toml
index 5910540..df23f52 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ dependencies = [
     "cloudpickle",
     "loky",
     "pluggy>=1.0.0",
-    "pytask>=0.4.0",
+    "pytask>=0.4.5",
     "rich"
 ]
 dynamic = ["version"]

From 6e8d1218b877ff536f66bbb6583525a9d6893f10 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Fri, 12 Jan 2024 16:11:48 +0100
Subject: [PATCH 5/6] to changes.

---
 CHANGES.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGES.md b/CHANGES.md
index 1210989..4a64b3e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -10,6 +10,7 @@ releases are available on [PyPI](https://pypi.org/project/pytask-parallel) and
 - {pull}`72` moves the project to `pyproject.toml`.
 - {pull}`75` updates the release strategy.
 - {pull}`79` add tests for Jupyter and fix parallelization with `PythonNode`s.
+- {pull}`82` fixes testing with pytask v0.4.5.
 
 ## 0.4.0 - 2023-10-07
 

From f85522f66690f9e1e894bccd403a46fcd5fd59d5 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Fri, 12 Jan 2024 16:20:22 +0100
Subject: [PATCH 6/6] Fix changes.

---
 CHANGES.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 370ca47..a2a5f88 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -5,7 +5,7 @@ chronological order. Releases follow [semantic versioning](https://semver.org/)
 releases are available on [PyPI](https://pypi.org/project/pytask-parallel) and
 [Anaconda.org](https://anaconda.org/conda-forge/pytask-parallel).
 
-## 0.4.1 - 2023-12-xx
+## 0.4.1 - 2024-01-12
 
 - {pull}`72` moves the project to `pyproject.toml`.
 - {pull}`75` updates the release strategy.