dag: add --outs option (#4739)

Fixes #4532
This commit is contained in:
Guro Bokum
2020-11-12 18:23:31 +09:00
committed by GitHub
parent 72eb4a5737
commit eeba567774
2 changed files with 47 additions and 4 deletions
+33 -4
View File
@@ -30,7 +30,7 @@ def _show_dot(G):
return dot_file.getvalue()
def _build(G, target=None, full=False):
def _build(G, target=None, full=False, outs=False):
import networkx as nx
from dvc.repo.graph import get_pipeline, get_pipelines
@@ -44,8 +44,25 @@ def _build(G, target=None, full=False):
else:
H = G
def _relabel(stage):
return stage.addressing
if outs:
G = nx.DiGraph()
for stage in H.nodes:
G.add_nodes_from(stage.outs)
for from_stage, to_stage in nx.edge_dfs(H):
G.add_edges_from(
[
(from_out, to_out)
for from_out in from_stage.outs
for to_out in to_stage.outs
]
)
H = G
def _relabel(node):
from dvc.stage import Stage
return node.addressing if isinstance(node, Stage) else str(node)
return nx.relabel_nodes(H, _relabel, copy=False)
@@ -64,7 +81,12 @@ class CmdDAG(CmdBase):
return 1
target = stages[0]
G = _build(self.repo.graph, target=target, full=self.args.full,)
G = _build(
self.repo.graph,
target=target,
full=self.args.full,
outs=self.args.outs,
)
if self.args.dot:
logger.info(_show_dot(G))
@@ -108,6 +130,13 @@ def add_parser(subparsers, parent_parser):
"showing DAG consisting only of ancestors."
),
)
dag_parser.add_argument(
"-o",
"--outs",
action="store_true",
default=False,
help="Print output files instead of stages.",
)
dag_parser.add_argument(
"target",
nargs="?",
+14
View File
@@ -58,6 +58,20 @@ def test_build_target(graph):
assert set(G.edges()) == {("3", "a.dvc"), ("3", "b.dvc")}
def test_build_target_with_outs(graph):
(stage,) = filter(
lambda s: hasattr(s, "name") and s.name == "3", graph.nodes()
)
G = _build(graph, target=stage, outs=True)
assert set(G.nodes()) == {"a", "b", "h", "i"}
assert set(G.edges()) == {
("h", "a"),
("h", "b"),
("i", "a"),
("i", "b"),
}
def test_build_full(graph):
(stage,) = filter(
lambda s: hasattr(s, "name") and s.name == "3", graph.nodes()