diff --git a/src/aiida_workgraph/decorator.py b/src/aiida_workgraph/decorator.py index 7627d553..a3eef220 100644 --- a/src/aiida_workgraph/decorator.py +++ b/src/aiida_workgraph/decorator.py @@ -178,8 +178,9 @@ def decorator_graph( def decorator(func) -> TaskHandle: from aiida_workgraph.tasks.graph_task import _build_graph_task_taskspec + from aiida_workgraph.task import GraphTaskHandle - handle = TaskHandle( + handle = GraphTaskHandle( _build_graph_task_taskspec( func, identifier=identifier, diff --git a/src/aiida_workgraph/task.py b/src/aiida_workgraph/task.py index 71cdcfa8..227544e5 100644 --- a/src/aiida_workgraph/task.py +++ b/src/aiida_workgraph/task.py @@ -11,7 +11,7 @@ from node_graph.task_spec import TaskSpec if TYPE_CHECKING: - pass + from aiida_workgraph.workgraph import WorkGraph class Task(GraphTask): @@ -231,15 +231,36 @@ def __call__(self, *args, **kwargs): return outputs - def run(self, /, *args, **kwargs): + +class GraphTaskHandle(TaskHandle): + """Handle for ``@task.graph`` callables. + + node-graph #150 made ``build`` graph-only by handle type (it moved from + ``BaseHandle`` to node-graph's ``GraphTaskHandle``). Mirror that here: + ``build``/``run``/``run_get_graph``/``submit`` all materialize a WorkGraph, + which only makes sense for graph specs — on plain task handles they only + ever raised. + + ``build`` delegates to node-graph's implementation unbound rather than + inheriting it, because node-graph's ``GraphTaskHandle.__init__`` takes + ``spec`` only, while ours must also wire ``get_current_graph`` and + ``graph_class=WorkGraph``. + """ + + def build(self, /, *args: Any, **kwargs: Any) -> WorkGraph: + from node_graph.task_spec import GraphTaskHandle as _NGGraphTaskHandle + + return _NGGraphTaskHandle.build(self, *args, **kwargs) + + def run(self, /, *args: Any, **kwargs: Any) -> Any: graph = self.build(*args, **kwargs) return graph.run() - def run_get_graph(self, /, *args, **kwargs): + def run_get_graph(self, /, *args: Any, **kwargs: Any) -> tuple[Any, WorkGraph]: graph = self.build(*args, **kwargs) return graph.run(), graph - def submit(self, /, *args, **kwargs): + def submit(self, /, *args: Any, **kwargs: Any) -> WorkGraph: graph = self.build(*args, **kwargs) graph.submit() return graph