Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 28 additions & 8 deletions src/_pytask/coiled_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,31 +3,51 @@
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Any
from typing import cast

if TYPE_CHECKING:
from collections.abc import Callable
from typing import Protocol

try:
from coiled.function import Function
except ImportError:

@dataclass
class Function:
cluster_kwargs: dict[str, Any]
environ: dict[str, Any]
_cluster_kwargs: dict[str, Any]
_environ: dict[str, Any] | None
_local: bool
_name: str
function: Callable[..., Any] | None
keepalive: str | None
keepalive: Any | None


__all__ = ["Function"]


if TYPE_CHECKING:

class _CoiledFunctionPrivateAttrs(Protocol):
_cluster_kwargs: dict[str, Any]
_environ: dict[str, Any] | None
_local: bool
_name: str
keepalive: Any | None


def _as_coiled_private_attrs(func: Function) -> _CoiledFunctionPrivateAttrs:
"""Cast to the private-attribute layout used by coiled's Function class."""
return cast("_CoiledFunctionPrivateAttrs", func)


def extract_coiled_function_kwargs(func: Function) -> dict[str, Any]:
"""Extract the kwargs for a coiled function."""
coiled_function = _as_coiled_private_attrs(func)
return {
"cluster_kwargs": func._cluster_kwargs, # ty: ignore[possibly-missing-attribute]
"keepalive": func.keepalive,
"environ": func._environ, # ty: ignore[possibly-missing-attribute]
"local": func._local, # ty: ignore[possibly-missing-attribute]
"name": func._name, # ty: ignore[possibly-missing-attribute]
"cluster_kwargs": coiled_function._cluster_kwargs,
"keepalive": coiled_function.keepalive,
"environ": coiled_function._environ,
"local": coiled_function._local,
"name": coiled_function._name,
}
4 changes: 2 additions & 2 deletions src/_pytask/path.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,10 +331,10 @@ def _import_module_using_spec(
break
else:
spec = importlib.util.spec_from_file_location(module_name, str(module_path))
if spec is not None:
if spec is not None and spec.loader is not None:
mod = importlib.util.module_from_spec(spec)
sys.modules[module_name] = mod
spec.loader.exec_module(mod) # ty: ignore[possibly-missing-attribute]
spec.loader.exec_module(mod)
_insert_missing_modules(sys.modules, module_name)
return mod

Expand Down
2 changes: 1 addition & 1 deletion src/_pytask/task_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ def wrapper(func: T) -> TaskDecorated[T]:

# When decorator is used without parentheses, call wrapper directly.
if is_task_function(name) and kwargs is None:
return wrapper(cast("T", name))
return wrapper(cast("T", name)) # ty: ignore[invalid-argument-type]
return wrapper


Expand Down
2 changes: 2 additions & 0 deletions tests/test_collect.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

from _pytask.collect import _find_shortest_uniquely_identifiable_name_for_tasks
from _pytask.collect import pytask_collect_node
from _pytask.node_protocols import PPathNode
from pytask import CollectionOutcome
from pytask import ExitCode
from pytask import NodeInfo
Expand Down Expand Up @@ -237,6 +238,7 @@ def test_pytask_collect_node_does_not_raise_error_if_path_is_not_normalized(
)
assert not record

assert isinstance(result, PPathNode)
assert str(result.path) == str(real_node)


Expand Down
3 changes: 2 additions & 1 deletion tests/test_dag_command.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import importlib
import os
import subprocess
import sys
Expand All @@ -12,7 +13,7 @@
from pytask import cli

try:
import pygraphviz # type: ignore[unresolved-import] # noqa: F401
importlib.import_module("pygraphviz")
except ImportError: # pragma: no cover
_IS_PYGRAPHVIZ_INSTALLED = False
else:
Expand Down
40 changes: 20 additions & 20 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.