Refactor exporting the DAG as a figure. (#185)

Author: tobiasraabe

docs/source/changes.rst

@@ -10,6 +10,8 @@ all releases are available on `PyPI <https://pypi.org/project/pytask>`_ and
 0.1.5 - 2022-xx-xx
 ------------------
 
+- :gh:`185` fix issues with drawing a graph and adds the ``--rank-direction`` to change
+  the direction of the DAG.
 - :gh:`186` enhance live displays by deactivating auto-refresh among other things.
 - :gh:`187` allows to enable and disable showing tracebacks and potentially different
   styles in the future with :confval:`show_traceback=True|False`.

src/_pytask/graph.py

@@ -4,6 +4,7 @@
 from typing import Any
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import TYPE_CHECKING
 
 import click
@@ -12,7 +13,6 @@
 from _pytask.compat import import_optional_dependency
 from _pytask.config import hookimpl
 from _pytask.console import console
-from _pytask.dag import descending_tasks
 from _pytask.exceptions import CollectionError
 from _pytask.exceptions import ConfigurationError
 from _pytask.exceptions import ResolvingDependenciesError
@@ -29,6 +29,13 @@
 if TYPE_CHECKING:
     from typing import NoReturn
 
+    if sys.version_info >= (3, 8):
+        from typing import Literal
+    else:
+        from typing_extensions import Literal
+
+    _RankDirection = Literal["TB", "LR", "BT", "RL"]
+
 
 @hookimpl(tryfirst=True)
 def pytask_extend_command_line_interface(cli: click.Group) -> None:
@@ -56,11 +63,33 @@ def pytask_parse_config(
         key="layout",
         default="dot",
     )
+    config["rank_direction"] = get_first_non_none_value(
+        config_from_cli,
+        config_from_file,
+        key="rank_direction",
+        default="TB",
+        callback=_rank_direction_callback,
+    )
+
+
+def _rank_direction_callback(
+    x: Optional["_RankDirection"],
+) -> Optional["_RankDirection"]:
+    """Validate the passed options for rank direction."""
+    if x in [None, "None", "none"]:
+        x = None
+    elif x in ["TB", "LR", "BT", "RL"]:
+        pass
+    else:
+        raise ValueError(
+            "'rank_direction' can only be one of ['TB', 'LR', 'BT', 'RL']."
+        )
+    return x
 
 
 _HELP_TEXT_LAYOUT: str = (
     "The layout determines the structure of the graph. Here you find an overview of "
-    "all available layouts: https://graphviz.org/#roadmap."
+    "all available layouts: https://graphviz.org/docs/layouts."
 )
 
 
@@ -70,9 +99,20 @@ def pytask_parse_config(
 )
 
 
+_HELP_TEXT_RANK_DIRECTION: str = (
+    "The direction of the directed graph. It can be ordered from top to bottom, TB, "
+    "left to right, LR, bottom to top, BT, or right to left, RL.  [default: TB]"
+)
+
+
 @click.command()
 @click.option("-l", "--layout", type=str, default=None, help=_HELP_TEXT_LAYOUT)
 @click.option("-o", "--output-path", type=str, default=None, help=_HELP_TEXT_OUTPUT)
+@click.option(
+    "--rank-direction",
+    type=click.Choice(["TB", "LR", "BT", "RL"]),
+    help=_HELP_TEXT_RANK_DIRECTION,
+)
 def dag(**config_from_cli: Any) -> "NoReturn":
     """Create a visualization of the project's DAG."""
     try:
@@ -181,10 +221,10 @@ def build_dag(config_from_cli: Dict[str, Any]) -> nx.DiGraph:
 def _refine_dag(session: Session) -> nx.DiGraph:
     """Refine the dag for plotting."""
     dag = _shorten_node_labels(session.dag, session.config["paths"])
-    dag = _add_root_node(dag)
     dag = _clean_dag(dag)
     dag = _style_dag(dag)
     dag = _escape_node_names_with_colons(dag)
+    dag.graph["graph"] = {"rankdir": session.config["rank_direction"]}
 
     return dag
 
@@ -239,21 +279,6 @@ def _shorten_node_labels(dag: nx.DiGraph, paths: List[Path]) -> nx.DiGraph:
     return dag
 
 
-def _add_root_node(dag: nx.DiGraph) -> nx.DiGraph:
-    """Add a root node to the graph to bind all starting nodes together."""
-    tasks_without_predecessor = [
-        name
-        for name in dag.nodes
-        if len(list(descending_tasks(name, dag))) == 0 and "task" in dag.nodes[name]
-    ]
-    if tasks_without_predecessor:
-        dag.add_node("root")
-        for name in tasks_without_predecessor:
-            dag.add_edge("root", name)
-
-    return dag
-
-
 def _clean_dag(dag: nx.DiGraph) -> nx.DiGraph:
     """Clean the DAG."""
     for node in dag.nodes:

src/_pytask/logging.py

@@ -104,7 +104,7 @@ def _show_traceback_callback(
     elif x in ["no", "yes"]:
         pass
     else:
-        raise ValueError("'show_traceback' can only be one of ['no', 'yes'")
+        raise ValueError("'show_traceback' can only be one of ['no', 'yes'].")
     return x