diff --git a/dvc/repo/imports.py b/dvc/repo/imports.py index 3dd35dfb8..a62072edf 100644 --- a/dvc/repo/imports.py +++ b/dvc/repo/imports.py @@ -44,7 +44,7 @@ def unfetched_view( return True - unfetched = index.targets_view(targets, filter_fn=need_fetch, **kwargs) + unfetched = index.targets_view(targets, stage_filter=need_fetch, **kwargs) return unfetched, changed_deps @@ -53,7 +53,7 @@ def partial_view( ) -> "IndexView": return index.targets_view( targets, - filter_fn=lambda s: s.is_partial_import, + stage_filter=lambda s: s.is_partial_import, **kwargs, ) diff --git a/dvc/repo/index.py b/dvc/repo/index.py index f3c2e1739..94fbe635f 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -151,24 +151,34 @@ class Index: def targets_view( self, targets: "TargetType", - filter_fn: Optional[Callable[["Stage"], bool]] = None, + stage_filter: Optional[Callable[["Stage"], bool]] = None, + outs_filter: Optional[Callable[["Output"], bool]] = None, **kwargs: Any, ) -> "IndexView": """Return read-only view of index for the specified targets. Args: targets: Targets to collect - filter_fn: Optional stage filter to be applied after collecting + stage_filter: Optional stage filter to be applied after collecting + targets. + outs_filter: Optional output filter to be applied after collecting targets. Additional kwargs will be passed into the stage collector. + + Note: + If both stage_filter and outs_filter are provided, stage_filter + will be applied first, and the resulting view will only contain + outputs from stages that matched stage_filter. Outputs from stages + that did not match will be excluded from the view (whether or not + the output would have matched outs_filter). """ stage_infos = [ stage_info for stage_info in self._collect_targets(targets, **kwargs) - if not filter_fn or filter_fn(stage_info.stage) + if not stage_filter or stage_filter(stage_info.stage) ] - return IndexView(self, stage_infos) + return IndexView(self, stage_infos, outs_filter=outs_filter) @property def outs(self) -> Iterator["Output"]: @@ -342,13 +352,17 @@ class IndexView: """Read-only view of Index.data using filtered stages.""" def __init__( # pylint: disable=redefined-outer-name - self, index: Index, stage_infos: Iterable["StageInfo"] + self, + index: Index, + stage_infos: Iterable["StageInfo"], + outs_filter: Optional[Callable[["Output"], bool]], ): self._index = index self._stage_infos = stage_infos # NOTE: stage_infos might have the same stage multiple times but with # different filter_info self._stages = list({stage for stage, _ in stage_infos}) + self._outs_filter = outs_filter def __len__(self) -> int: return len(self._stages) @@ -381,7 +395,8 @@ class IndexView: outs = set() for stage, filter_info in self._stage_infos: for out in stage.filter_outs(filter_info): - outs.add(out) + if not self._outs_filter or self._outs_filter(out): + outs.add(out) yield from outs @cached_property @@ -391,6 +406,8 @@ class IndexView: prefixes: Dict[str, Set["DataIndexKey"]] = defaultdict(set) for stage, filter_info in self._stage_infos: for out in stage.filter_outs(filter_info): + if self._outs_filter and not self._outs_filter(out): + continue workspace, key = out.index_key if filter_info and out.fs.path.isin(filter_info, out.fs_path): key = ( diff --git a/tests/func/test_repo_index.py b/tests/func/test_repo_index.py index 516cd37f1..0056c97aa 100644 --- a/tests/func/test_repo_index.py +++ b/tests/func/test_repo_index.py @@ -303,3 +303,62 @@ def test_getitem(tmp_dir, dvc, run_copy): with pytest.raises(KeyError): _ = index["no-valid-stage-name"] + + +def test_view_stage_filter(tmp_dir, scm, dvc, run_copy): + (stage1,) = tmp_dir.dvc_gen("foo", "foo") + stage2 = run_copy("foo", "bar", name="copy-foo-bar") + tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add") + index = Index(dvc) + + view = index.targets_view(None) + assert set(view.stages) == {stage1, stage2} + assert {out.fs_path for out in view.outs} == { + out.fs_path for out in (stage1.outs + stage2.outs) + } + + view = index.targets_view( + None, stage_filter=lambda s: getattr(s, "name", "").startswith("copy") + ) + assert set(view.stages) == {stage2} + assert {out.fs_path for out in view.outs} == { + out.fs_path for out in stage2.outs + } + + +def test_view_outs_filter(tmp_dir, scm, dvc, run_copy): + (stage1,) = tmp_dir.dvc_gen("foo", "foo") + stage2 = run_copy("foo", "bar", name="copy-foo-bar") + tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add") + index = Index(dvc) + + view = index.targets_view(None, outs_filter=lambda o: o.def_path == "foo") + assert set(view.stages) == {stage1, stage2} + assert {out.fs_path for out in view.outs} == { + out.fs_path for out in stage1.outs + } + + +def test_view_combined_filter(tmp_dir, scm, dvc, run_copy): + (stage1,) = tmp_dir.dvc_gen("foo", "foo") + stage2 = run_copy("foo", "bar", name="copy-foo-bar") + tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add") + index = Index(dvc) + + view = index.targets_view( + None, + stage_filter=lambda s: getattr(s, "name", "").startswith("copy"), + outs_filter=lambda o: o.def_path == "foo", + ) + assert set(view.stages) == {stage2} + assert set(view.outs) == set() + + view = index.targets_view( + None, + stage_filter=lambda s: getattr(s, "name", "").startswith("copy"), + outs_filter=lambda o: o.def_path == "bar", + ) + assert set(view.stages) == {stage2} + assert {out.fs_path for out in view.outs} == { + out.fs_path for out in stage2.outs + }