index: fix granular dir output view filtering

This commit is contained in:
Peter Rowlands
2022-11-08 19:28:10 +09:00
committed by Peter Rowlands (변기호)
parent 99e7f52bbd
commit 628b992442
2 changed files with 52 additions and 22 deletions
+29 -22
View File
@@ -7,8 +7,10 @@ from typing import (
Iterable,
Iterator,
List,
NamedTuple,
Optional,
Set,
Tuple,
)
from funcy import cached_property, nullcontext
@@ -348,6 +350,11 @@ class Index:
return dict_md5(self.dumpd())
class _DataPrefixes(NamedTuple):
explicit: Set["DataIndexKey"]
recursive: Set["DataIndexKey"]
class IndexView:
"""Read-only view of Index.data using filtered stages."""
@@ -391,33 +398,32 @@ class IndexView:
yield from stage.deps
@property
def outs(self) -> Iterator["Output"]:
outs = set()
def _filtered_outs(self) -> Iterator[Tuple["Output", Optional[str]]]:
for stage, filter_info in self._stage_infos:
for out in stage.filter_outs(filter_info):
if not self._outs_filter or self._outs_filter(out):
outs.add(out)
yield from outs
yield out, filter_info
@property
def outs(self) -> Iterator["Output"]:
yield from {out for (out, _) in self._filtered_outs}
@cached_property
def _data_prefixes(self) -> Dict[str, Set["DataIndexKey"]]:
def _data_prefixes(self) -> Dict[str, "_DataPrefixes"]:
from collections import defaultdict
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 = (
*key,
out.fs.path.relparts(filter_info, out.fs_path),
)
prefixes[workspace].add(key)
prefixes[workspace].update(
key[:i] for i in range(len(key), 0, -1)
)
prefixes: Dict[str, "_DataPrefixes"] = defaultdict(
lambda: _DataPrefixes(set(), set())
)
for out, filter_info in self._filtered_outs:
workspace, key = out.index_key
if filter_info and out.fs.path.isin(filter_info, out.fs_path):
key = key + out.fs.path.relparts(filter_info, out.fs_path)
if out.meta.isdir:
prefixes[workspace].recursive.add(key)
prefixes[workspace].explicit.update(
key[:i] for i in range(len(key), 0, -1)
)
return prefixes
@cached_property
@@ -428,9 +434,10 @@ class IndexView:
def key_filter(workspace: str, key: "DataIndexKey"):
try:
return key in self._data_prefixes[workspace] or any(
prefixes = self._data_prefixes[workspace]
return key in prefixes.explicit or any(
key[: len(prefix)] == prefix
for prefix in self._data_prefixes[workspace]
for prefix in prefixes.recursive
)
except KeyError:
return False
+23
View File
@@ -305,6 +305,29 @@ def test_getitem(tmp_dir, dvc, run_copy):
_ = index["no-valid-stage-name"]
def test_view_granular_dir(tmp_dir, scm, dvc, run_copy):
tmp_dir.dvc_gen(
{"dir": {"subdir": {"in-subdir": "in-subdir"}, "in-dir": "in-dir"}},
commit="init",
)
index = Index(dvc)
# view should include the specific target, parent dirs, and children
# view should exclude any siblings of the target
view = index.targets_view("dir/subdir")
data_index = view.data["repo"]
assert ("dir",) in data_index
assert (
"dir",
"subdir",
) in data_index
assert ("dir", "subdir", "in-subdir") in data_index
assert (
"dir",
"in-dir",
) not in data_index
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")