From c29e721823f5cd5bdedcfd3c164e1996a33e594b Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Tue, 13 Sep 2022 17:25:12 +0900 Subject: [PATCH] output: support loading/dumping dir trees from metadata --- dvc/dependency/param.py | 2 +- dvc/dependency/repo.py | 2 +- dvc/dvcfile.py | 2 +- dvc/output.py | 33 +++++++++++++++++++++++++++++++-- dvc/stage/__init__.py | 4 ++-- dvc/stage/serialize.py | 4 ++-- dvc/stage/utils.py | 6 +++--- 7 files changed, 41 insertions(+), 12 deletions(-) diff --git a/dvc/dependency/param.py b/dvc/dependency/param.py index f5ce0a5fb..10257f3bd 100644 --- a/dvc/dependency/param.py +++ b/dvc/dependency/param.py @@ -47,7 +47,7 @@ class ParamsDependency(Dependency): super().__init__(stage, path, repo=repo) self.hash_info = hash_info - def dumpd(self): + def dumpd(self, **kwargs): ret = super().dumpd() if not self.hash_info: ret[self.PARAM_PARAMS] = self.params or {} diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index 5a95da192..df75a8191 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -59,7 +59,7 @@ class RepoDependency(Dependency): def save(self): pass - def dumpd(self): + def dumpd(self, **kwargs): return {self.PARAM_PATH: self.def_path, self.PARAM_REPO: self.def_repo} def download(self, to, jobs=None): diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index fb860d24d..9d2098245 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -198,7 +198,7 @@ class SingleStageFile(FileMixin): if self.verify: check_dvcfile_path(self.repo, self.path) logger.debug("Saving information to '%s'.", relpath(self.path)) - dump_yaml(self.path, serialize.to_single_stage_file(stage)) + dump_yaml(self.path, serialize.to_single_stage_file(stage, **kwargs)) self.repo.scm_context.track_file(self.relpath) def remove_stage(self, stage): # pylint: disable=unused-argument diff --git a/dvc/output.py b/dvc/output.py index 67023922e..bd55e0750 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -85,6 +85,7 @@ def loadd_from(stage, d_list): live = d.pop(Output.PARAM_LIVE, False) remote = d.pop(Output.PARAM_REMOTE, None) annot = {field: d.pop(field, None) for field in ANNOTATION_FIELDS} + files = d.pop(Output.PARAM_FILES, None) ret.append( _get( stage, @@ -98,6 +99,7 @@ def loadd_from(stage, d_list): live=live, remote=remote, **annot, + files=files, ) ) return ret @@ -243,6 +245,7 @@ class Output: PARAM_PATH = "path" PARAM_CACHE = "cache" PARAM_CHECKPOINT = "checkpoint" + PARAM_FILES = "files" PARAM_METRIC = "metric" PARAM_METRIC_TYPE = "type" PARAM_METRIC_XPATH = "xpath" @@ -292,6 +295,7 @@ class Output: remote=None, repo=None, fs_config=None, + files=None, ): self.annot = Annotation( desc=desc, type=type, labels=labels or [], meta=meta or {} @@ -340,6 +344,7 @@ class Output: # should be absolute and don't contain remote:// refs. self.stage = stage self.meta = meta + self._dump_files = files is not None self.use_cache = False if self.IS_DEPENDENCY else cache self.metric = False if self.IS_DEPENDENCY else metric self.plot = False if self.IS_DEPENDENCY else plot @@ -367,6 +372,12 @@ class Output: name=self.hash_name, value=getattr(self.meta, self.hash_name, None), ) + if self.hash_info and self.hash_info.isdir: + self.meta.isdir = True + if files: + tree = Tree.from_list(files, hash_name=self.hash_name) + tree.digest() + self.obj = tree def _parse_path(self, fs, fs_path): parsed = urlparse(self.def_path) @@ -708,8 +719,10 @@ class Output: ) return checkout_obj - def dumpd(self): - ret = {**self.hash_info.to_dict(), **self.meta.to_dict()} + def dumpd(self, **kwargs): + meta = self.meta.to_dict() + meta.pop("isdir", None) + ret = {**self.hash_info.to_dict(), **meta} if self.is_in_repo: path = self.fs.path.as_posix( @@ -752,6 +765,15 @@ class Output: if self.remote: ret[self.PARAM_REMOTE] = self.remote + if ( + self.use_cache + and self.is_in_repo + and self.hash_info.isdir + and (kwargs.get("with_files") or self._dump_files) + ): + assert self.obj + ret[self.PARAM_FILES] = self.obj.as_list(with_meta=True) + return ret def verify_metric(self): @@ -1155,10 +1177,17 @@ ARTIFACT_SCHEMA = { Output.PARAM_CHECKPOINT: bool, } +DIR_FILES_SCHEMA: Dict[str, Any] = { + **CHECKSUMS_SCHEMA, + **META_SCHEMA, + Required(Tree.PARAM_RELPATH): str, +} + SCHEMA = { **ARTIFACT_SCHEMA, **ANNOTATION_SCHEMA, Output.PARAM_CACHE: bool, Output.PARAM_METRIC: Output.METRIC_SCHEMA, Output.PARAM_REMOTE: str, + Output.PARAM_FILES: [DIR_FILES_SCHEMA], } diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 4dd5a923c..6d19d4f93 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -461,8 +461,8 @@ class Stage(params.StageParams): def reload(self): return self.dvcfile.stage - def dumpd(self): - return get_dump(self) + def dumpd(self, **kwargs): + return get_dump(self, **kwargs) def compute_md5(self): # `dvc add`ed files don't need stage md5 diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index f918ad956..b3da928eb 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -188,8 +188,8 @@ def to_lockfile(stage: "PipelineStage") -> dict: return {stage.name: to_single_stage_lockfile(stage)} -def to_single_stage_file(stage: "Stage"): - state = stage.dumpd() +def to_single_stage_file(stage: "Stage", **kwargs): + state = stage.dumpd(**kwargs) # When we load a stage we parse yaml with a fast parser, which strips # off all the comments and formatting. To retain those on update we do diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py index 6db4694a8..bb287e843 100644 --- a/dvc/stage/utils.py +++ b/dvc/stage/utils.py @@ -226,7 +226,7 @@ def resolve_paths(fs, path, wdir=None): return path, wdir -def get_dump(stage): +def get_dump(stage, **kwargs): return { key: value for key, value in { @@ -235,8 +235,8 @@ def get_dump(stage): stage.PARAM_CMD: stage.cmd, stage.PARAM_WDIR: resolve_wdir(stage.wdir, stage.path), stage.PARAM_FROZEN: stage.frozen, - stage.PARAM_DEPS: [d.dumpd() for d in stage.deps], - stage.PARAM_OUTS: [o.dumpd() for o in stage.outs], + stage.PARAM_DEPS: [d.dumpd(**kwargs) for d in stage.deps], + stage.PARAM_OUTS: [o.dumpd(**kwargs) for o in stage.outs], stage.PARAM_ALWAYS_CHANGED: stage.always_changed, stage.PARAM_META: stage.meta, }.items()