Fix.

Author: tobiasraabe

src/_pytask/dag.py

@@ -1,4 +1,4 @@
-"""Contains code related to resolving dependencies."""
+"""Contains code related to the DAG."""
 
 from __future__ import annotations
 
@@ -34,18 +34,13 @@
     from _pytask.session import Session
 
 
-__all__ = ["create_dag"]
+__all__ = ["create_dag", "create_dag_from_session"]
 
 
 def create_dag(session: Session) -> nx.DiGraph:
     """Create a directed acyclic graph (DAG) for the workflow."""
     try:
-        dag = _create_dag(tasks=session.tasks)
-        _check_if_dag_has_cycles(dag)
-        _check_if_tasks_have_the_same_products(dag, session.config["paths"])
-        _modify_dag(session=session, dag=dag)
-        select_tasks_by_marks_and_expressions(session=session, dag=dag)
-
+        dag = create_dag_from_session(session)
     except Exception:  # noqa: BLE001
         report = DagReport.from_exception(sys.exc_info())
         _log_dag(report=report)
@@ -55,7 +50,17 @@ def create_dag(session: Session) -> nx.DiGraph:
     return dag
 
 
-def _create_dag(tasks: list[PTask]) -> nx.DiGraph:
+def create_dag_from_session(session: Session) -> nx.DiGraph:
+    """Create a DAG from a session."""
+    dag = _create_dag_from_tasks(tasks=session.tasks)
+    _check_if_dag_has_cycles(dag)
+    _check_if_tasks_have_the_same_products(dag, session.config["paths"])
+    dag = _modify_dag(session=session, dag=dag)
+    select_tasks_by_marks_and_expressions(session=session, dag=dag)
+    return dag
+
+
+def _create_dag_from_tasks(tasks: list[PTask]) -> nx.DiGraph:
     """Create the DAG from tasks, dependencies and products."""
 
     def _add_dependency(
@@ -98,7 +103,7 @@ def _add_product(
     return dag
 
 
-def _modify_dag(session: Session, dag: nx.DiGraph) -> None:
+def _modify_dag(session: Session, dag: nx.DiGraph) -> nx.DiGraph:
     """Create dependencies between tasks when using ``@task(after=...)``."""
     temporary_id_to_task = {
         task.attributes["collection_id"]: task
@@ -119,6 +124,7 @@ def _modify_dag(session: Session, dag: nx.DiGraph) -> None:
             for signature in signatures:
                 for successor in dag.successors(signature):
                     dag.add_edge(successor, task.signature)
+    return dag
 
 
 def _check_if_dag_has_cycles(dag: nx.DiGraph) -> None:

src/_pytask/provisional_utils.py

@@ -1,3 +1,5 @@
+"""Contains utilities related to provisional nodes and task generators."""
+
 from __future__ import annotations
 
 import sys
@@ -6,12 +8,8 @@
 from typing import Any
 
 from _pytask.collect_utils import collect_dependency
-from _pytask.dag import _check_if_dag_has_cycles
-from _pytask.dag import _check_if_tasks_have_the_same_products
-from _pytask.dag import _create_dag
-from _pytask.dag import _modify_dag
+from _pytask.dag import create_dag_from_session
 from _pytask.dag_utils import TopologicalSorter
-from _pytask.mark import select_tasks_by_marks_and_expressions
 from _pytask.models import NodeInfo
 from _pytask.node_protocols import PNode
 from _pytask.node_protocols import PProvisionalNode
@@ -72,14 +70,14 @@ def collect_provisional_nodes(
 
 
 def recreate_dag(session: Session, task: PTask) -> None:
-    """Recreate the DAG."""
+    """Recreate the DAG when provisional nodes are resolved.
+
+    If the DAG resolution fails, the error is attached as an execution report since
+    there is not better mechanic yet to display the error.
+
+    """
     try:
-        dag = _create_dag(tasks=session.tasks)
-        _check_if_dag_has_cycles(dag)
-        _check_if_tasks_have_the_same_products(dag, session.config["paths"])
-        _modify_dag(session=session, dag=dag)
-        select_tasks_by_marks_and_expressions(session=session, dag=dag)
-        session.dag = dag
+        session.dag = create_dag_from_session(session)
         session.scheduler = TopologicalSorter.from_dag_and_sorter(
             session.dag, session.scheduler
         )

tests/test_dag.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 import pytest
-from _pytask.dag import _create_dag
+from _pytask.dag import _create_dag_from_tasks
 from pytask import ExitCode
 from pytask import PathNode
 from pytask import Task
@@ -26,7 +26,7 @@ def test_create_dag():
             1: PathNode.from_path(root / "node_2"),
         },
     )
-    dag = _create_dag(tasks=[task])
+    dag = _create_dag_from_tasks(tasks=[task])
 
     for signature in (
         "90bb899a1b60da28ff70352cfb9f34a8bed485597c7f40eed9bd4c6449147525",