diff --git a/dvc/repo/index.py b/dvc/repo/index.py index 36397007d..bfc725d65 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -26,7 +26,7 @@ if TYPE_CHECKING: from dvc.stage import Stage from dvc.types import StrPath, TargetType from dvc_data.hashfile.hash_info import HashInfo - from dvc_data.index import DataIndex + from dvc_data.index import DataIndex, DataIndexKey, DataIndexView from dvc_objects.db import ObjectDB @@ -148,6 +148,28 @@ class Index: ) return self._collected_targets[targets_hash] + def targets_view( + self, + targets: "TargetType", + filter_fn: Optional[Callable[["Stage"], 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 + targets. + + Additional kwargs will be passed into the stage collector. + """ + stage_infos = [ + stage_info + for stage_info in self._collect_targets(targets, **kwargs) + if not filter_fn or filter_fn(stage_info.stage) + ] + return IndexView(self, stage_infos) + @property def outs(self) -> Iterator["Output"]: for stage in self: @@ -339,6 +361,75 @@ class Index: return dict_md5(self.dumpd()) +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 + self._stages: Dict["Stage", Optional[str]] = dict(stage_infos) + + def __len__(self) -> int: + return len(self._stages) + + def __contains__(self, stage: "Stage") -> bool: + # as we are keeping stages inside a list, it might be slower. + return stage in self._stages + + def __iter__(self) -> Iterator["Stage"]: + yield from self._stages + + def __getitem__(self, item: str) -> "Stage": + """Get a stage by its addressing attribute.""" + for stage in self: + if stage.addressing == item: + return stage + raise KeyError(f"{item} - available stages are {self.stages}") + + @property + def stages(self): + return list(self._stages) + + @property + def deps(self) -> Iterator["Dependency"]: + for stage in self: + yield from stage.deps + + @property + def outs(self) -> Iterator["Output"]: + for stage, filter_info in self._stages.items(): + yield from stage.filter_outs(filter_info) + + @cached_property + def _data_prefixes(self) -> Dict[str, Set["DataIndexKey"]]: + from collections import defaultdict + + prefixes: Dict[str, Set["DataIndexKey"]] = defaultdict(set) + for out in self.outs: + workspace, key = out.index_key + prefixes[workspace].add(key) + prefixes[workspace].update(key[:i] for i in range(len(key), 0, -1)) + return prefixes + + @cached_property + def data(self) -> "Dict[str, DataIndexView]": + from functools import partial + + from dvc_data.index import view + + def key_filter(workspace: str, key: "DataIndexKey"): + try: + return key in self._data_prefixes[workspace] + except KeyError: + return False + + return { + workspace: view(index, partial(key_filter, workspace)) + for workspace, index in self._index.data.items() + } + + if __name__ == "__main__": import logging