output: support loading/dumping dir trees from metadata

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 {}

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):

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

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],
 }

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

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

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()