Stop unwrapping coiled functions.

Author: tobiasraabe

docs/source/changes.md

@@ -39,6 +39,7 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 - {pull}`593` recreate `PythonNode`s every run since they carry the `_NoDefault` enum as
   the value whose state is `None`.
 - {pull}`594` publishes `NodeLoadError`.
+- {pull}`595` stops unwrapping task functions until a `coiled.function.Function`.
 
 ## 0.4.7 - 2024-03-19
 

src/_pytask/collect.py

@@ -46,6 +46,7 @@
 from _pytask.reports import CollectionReport
 from _pytask.shared import find_duplicates
 from _pytask.shared import to_list
+from _pytask.shared import unwrap_task_function
 from _pytask.task_utils import COLLECTED_TASKS
 from _pytask.task_utils import task as task_decorator
 from _pytask.typing import is_task_function
@@ -317,9 +318,7 @@ def pytask_collect_task(
             obj.pytask_meta.is_generator if hasattr(obj, "pytask_meta") else False
         )
 
-        # Get the underlying function to avoid having different states of the function,
-        # e.g. due to pytask_meta, in different layers of the wrapping.
-        unwrapped = inspect.unwrap(obj)
+        unwrapped = unwrap_task_function(obj)
 
         if path is None:
             return TaskWithoutPath(

src/_pytask/shared.py

@@ -3,13 +3,16 @@
 from __future__ import annotations
 
 import glob
+import inspect
 from pathlib import Path
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Callable
 from typing import Iterable
 from typing import Sequence
 
 import click
+from attrs import define
 
 from _pytask.console import format_node_name
 from _pytask.console import format_task_name
@@ -23,13 +26,26 @@
     import networkx as nx
 
 
+try:
+    from coiled.function import Function as CoiledFunction
+except ImportError:
+
+    @define
+    class CoiledFunction:  # type: ignore[no-redef]
+        cluster_kwargs: dict[str, Any]
+        environ: dict[str, Any]
+        function: Callable[..., Any] | None
+        keepalive: str | None
+
+
 __all__ = [
     "convert_to_enum",
     "find_duplicates",
     "parse_markers",
     "parse_paths",
     "reduce_names_of_multiple_nodes",
     "to_list",
+    "unwrap_task_function",
 ]
 
 
@@ -146,3 +162,13 @@ def convert_to_enum(value: Any, enum: type[Enum]) -> Enum:
         values = [e.value for e in enum]
         msg = f"Value {value!r} is not a valid {enum!r}. Valid values are {values}."
         raise ValueError(msg) from None
+
+
+def unwrap_task_function(obj: Any) -> Callable[..., Any]:
+    """Unwrap a task function.
+
+    Get the underlying function to avoid having different states of the function, e.g.
+    due to pytask_meta, in different layers of the wrapping.
+
+    """
+    return inspect.unwrap(obj, stop=lambda x: isinstance(x, CoiledFunction))

src/_pytask/task_utils.py

@@ -16,6 +16,7 @@
 from _pytask.mark import Mark
 from _pytask.models import CollectionMetadata
 from _pytask.shared import find_duplicates
+from _pytask.shared import unwrap_task_function
 from _pytask.typing import is_task_function
 
 if TYPE_CHECKING:
@@ -117,7 +118,7 @@ def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
                 )
                 raise ValueError(msg)
 
-        unwrapped = inspect.unwrap(func)
+        unwrapped = unwrap_task_function(func)
 
         # We do not allow builtins as functions because we would need to use
         # ``inspect.stack`` to infer their caller location and they are unable to carry
@@ -145,7 +146,7 @@ def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
             unwrapped.pytask_meta.produces = produces
             unwrapped.pytask_meta.after = parsed_after
         else:
-            unwrapped.pytask_meta = CollectionMetadata(
+            unwrapped.pytask_meta = CollectionMetadata(  # type: ignore[attr-defined]
                 after=parsed_after,
                 is_generator=is_generator,
                 id_=id,

tests/test_shared.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import functools
 import textwrap
 from contextlib import ExitStack as does_not_raise  # noqa: N813
 
 import pytest
 from _pytask.shared import convert_to_enum
 from _pytask.shared import find_duplicates
+from _pytask.shared import unwrap_task_function
 from pytask import ExitCode
 from pytask import ShowCapture
 from pytask import build
@@ -49,3 +51,33 @@ def test_convert_to_enum(value, enum, expectation, expected):
     with expectation:
         result = convert_to_enum(value, enum)
         assert result == expected
+
+
+@pytest.mark.unit()
+def test_unwrap_task_function():
+    def task():
+        pass
+
+    # partialed functions are only unwrapped after wraps.
+    partialed = functools.wraps(task)(functools.partial(task))
+    assert unwrap_task_function(partialed) is task
+
+    partialed = functools.partial(task)
+    assert unwrap_task_function(partialed) is partialed
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper():
+            return func()
+
+        return wrapper
+
+    decorated = decorator(task)
+    assert unwrap_task_function(decorated) is task
+
+    from _pytask.shared import CoiledFunction
+
+    coiled_function = functools.wraps(task)(
+        CoiledFunction(function=task, cluster_kwargs={}, environ={}, keepalive=None)
+    )
+    assert unwrap_task_function(coiled_function) is coiled_function