Simplify the tree_map code for generating the DAG. (#447)

docs/source/changes.md

@@ -5,13 +5,14 @@ chronological order. Releases follow [semantic versioning](https://semver.org/)
 releases are available on [PyPI](https://pypi.org/project/pytask) and
 [Anaconda.org](https://anaconda.org/conda-forge/pytask).
 
-## 0.4.1 - 2023-10-xx
+## 0.4.1 - 2023-10-11
 
 - {pull}`443` ensures that `PythonNode.name` is always unique by only handling it
   internally.
 - {pull}`444` moves all content of `setup.cfg` to `pyproject.toml`.
 - {pull}`446` refactors `create_name_of_python_node` and fixes `PythonNode`s as returns.
-- {pull}`447` fixes handling multiple product annotations of a task.
+- {pull}`447` simplifies the `tree_map` code while generating the DAG.
+- {pull}`448` fixes handling multiple product annotations of a task.
 
 ## 0.4.0 - 2023-10-07
 

src/_pytask/dag.py

@@ -76,16 +76,30 @@ def pytask_dag(session: Session) -> bool | None:
 @hookimpl
 def pytask_dag_create_dag(tasks: list[PTask]) -> nx.DiGraph:
     """Create the DAG from tasks, dependencies and products."""
+
+    def _add_dependency(dag: nx.DiGraph, task: PTask, node: PNode) -> None:
+        """Add a dependency to the DAG."""
+        dag.add_node(node.name, node=node)
+        dag.add_edge(node.name, task.name)
+
+        # If a node is a PythonNode wrapped in another PythonNode, it is a product from
+        # another task that is a dependency in the current task. Thus, draw an edge
+        # connecting the two nodes.
+        if isinstance(node, PythonNode) and isinstance(node.value, PythonNode):
+            dag.add_edge(node.value.name, node.name)
+
+    def _add_product(dag: nx.DiGraph, task: PTask, node: PNode) -> None:
+        """Add a product to the DAG."""
+        dag.add_node(node.name, node=node)
+        dag.add_edge(task.name, node.name)
+
     dag = nx.DiGraph()
 
     for task in tasks:
         dag.add_node(task.name, task=task)
 
-        tree_map(lambda x: dag.add_node(x.name, node=x), task.depends_on)
-        tree_map(lambda x: dag.add_edge(x.name, task.name), task.depends_on)
-
-        tree_map(lambda x: dag.add_node(x.name, node=x), task.produces)
-        tree_map(lambda x: dag.add_edge(task.name, x.name), task.produces)
+        tree_map(lambda x: _add_dependency(dag, task, x), task.depends_on)
+        tree_map(lambda x: _add_product(dag, task, x), task.produces)
 
         # If a node is a PythonNode wrapped in another PythonNode, it is a product from
         # another task that is a dependency in the current task. Thus, draw an edge