diff --git a/dvc/commands/dag.py b/dvc/commands/dag.py index 03cc96dad..8d4f87bfe 100644 --- a/dvc/commands/dag.py +++ b/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", diff --git a/tests/unit/command/test_compat_flag.py b/tests/unit/command/test_compat_flag.py index 645096b57..996121bbf 100644 --- a/tests/unit/command/test_compat_flag.py +++ b/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"), diff --git a/tests/unit/command/test_dag.py b/tests/unit/command/test_dag.py index 6d6d291a1..7a23d1167 100644 --- a/tests/unit/command/test_dag.py +++ b/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", + "```", + ]