dag: Add --mermaid option.

dvc/commands/dag.py

@@ -28,6 +28,32 @@ def _show_dot(G):
     return dot_file.getvalue()
 
 
+def _show_mermaid(G, markdown: bool = False):
+    from dvc.repo.graph import get_pipelines
+
+    pipelines = get_pipelines(G)
+
+    graph = "flowchart TD"
+
+    total_nodes = 0
+    for pipeline in pipelines:
+        node_ids = {}
+        nodes = sorted(str(x) for x in pipeline.nodes)
+        for node in nodes:
+            total_nodes += 1
+            node_id = f"node{total_nodes}"
+            graph += f"\n\t{node_id}[{node}]"
+            node_ids[node] = node_id
+        edges = sorted((str(a), str(b)) for b, a in pipeline.edges)
+        for a, b in edges:
+            graph += f"\n\t{node_ids[str(a)]}-->{node_ids[str(b)]}"
+
+    if markdown:
+        return f"```mermaid\n{graph}\n```"
+
+    return graph
+
+
 def _collect_targets(repo, target, outs):
     if not target:
         return []
@@ -105,6 +131,8 @@ class CmdDAG(CmdBase):
 
         if self.args.dot:
             ui.write(_show_dot(G))
+        elif self.args.mermaid or self.args.markdown:
+            ui.write(_show_mermaid(G, self.args.markdown))
         else:
             with ui.pager():
                 ui.write(_show_ascii(G))
@@ -127,6 +155,20 @@ def add_parser(subparsers, parent_parser):
         default=False,
         help="Print DAG with .dot format.",
     )
+    dag_parser.add_argument(
+        "--mermaid",
+        action="store_true",
+        default=False,
+        help="Print DAG with mermaid format.",
+    )
+    dag_parser.add_argument(
+        "--md",
+        "--show-md",
+        action="store_true",
+        default=False,
+        dest="markdown",
+        help="Print DAG with mermaid format wrapped in Markdown block.",
+    )
     dag_parser.add_argument(
         "--full",
         action="store_true",

tests/unit/command/test_compat_flag.py

@@ -14,6 +14,7 @@ def _id_gen(val) -> str:
 @pytest.mark.parametrize(
     "args, key",
     [
+        (["dag", "--show-md"], "markdown"),
         (["diff", "--show-json"], "json"),
         (["diff", "--show-md"], "markdown"),
         (["experiments", "diff", "--show-json"], "json"),

tests/unit/command/test_dag.py

@@ -1,12 +1,28 @@
 import networkx as nx
 import pytest
 
-from dvc.cli import parse_args
-from dvc.commands.dag import CmdDAG, _build, _show_ascii, _show_dot
+from dvc.cli import main, parse_args
+from dvc.commands.dag import (
+    CmdDAG,
+    _build,
+    _show_ascii,
+    _show_dot,
+    _show_mermaid,
+)
 
 
-@pytest.mark.parametrize("fmt", [None, "--dot"])
-def test_dag(tmp_dir, dvc, mocker, fmt):
+@pytest.mark.parametrize(
+    "fmt, formatter",
+    [
+        (None, "_show_ascii"),
+        ("--dot", "_show_dot"),
+        ("--mermaid", "_show_mermaid"),
+        ("--md", "_show_mermaid"),
+    ],
+)
+def test_dag(tmp_dir, dvc, mocker, fmt, formatter):
+    from dvc.commands import dag
+
     tmp_dir.dvc_gen("foo", "foo")
 
     args = ["dag", "--full", "foo.dvc"]
@@ -15,12 +31,16 @@ def test_dag(tmp_dir, dvc, mocker, fmt):
     cli_args = parse_args(args)
     assert cli_args.func == CmdDAG
 
+    fmt_func = mocker.spy(dag, formatter)
+
     cmd = cli_args.func(cli_args)
 
     mocker.patch("dvc.commands.dag._build", return_value=dvc.index.graph)
 
     assert cmd.run() == 0
 
+    assert fmt_func.called
+
 
 @pytest.fixture
 def repo(tmp_dir, dvc):
@@ -131,3 +151,49 @@ def test_show_dot(repo):
         "\"stage: '4'\" -> \"stage: '3'\";\n"
         "}\n"
     )
+
+
+def test_show_mermaid(repo):
+    assert [
+        line.rstrip() for line in _show_mermaid(repo.index.graph).splitlines()
+    ] == [
+        "flowchart TD",
+        "\tnode1[stage: '1']",
+        "\tnode2[stage: '2']",
+        "\tnode3[stage: '3']",
+        "\tnode4[stage: '4']",
+        "\tnode5[stage: 'a.dvc']",
+        "\tnode6[stage: 'b.dvc']",
+        "\tnode3-->node4",
+        "\tnode5-->node1",
+        "\tnode5-->node3",
+        "\tnode5-->node4",
+        "\tnode6-->node2",
+        "\tnode6-->node3",
+    ]
+
+
+def test_show_mermaid_markdown(repo, dvc, capsys, mocker):
+    mocker.patch("dvc.commands.dag._build", return_value=dvc.index.graph)
+
+    capsys.readouterr()
+    assert main(["dag", "--md"]) == 0
+    assert [
+        line.rstrip() for line in capsys.readouterr().out.splitlines()
+    ] == [
+        "```mermaid",
+        "flowchart TD",
+        "\tnode1[stage: '1']",
+        "\tnode2[stage: '2']",
+        "\tnode3[stage: '3']",
+        "\tnode4[stage: '4']",
+        "\tnode5[stage: 'a.dvc']",
+        "\tnode6[stage: 'b.dvc']",
+        "\tnode3-->node4",
+        "\tnode5-->node1",
+        "\tnode5-->node3",
+        "\tnode5-->node4",
+        "\tnode6-->node2",
+        "\tnode6-->node3",
+        "```",
+    ]