From 0e0d46e2325d3718a3063141f52ccac09ca6bddb Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Thu, 15 Sep 2022 15:47:34 +0900
Subject: [PATCH] output: lazy load tree from files list in get_obj

---
 dvc/output.py | 46 ++++++++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/dvc/output.py b/dvc/output.py
index bd55e0750..904d21713 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -1,7 +1,7 @@
 import logging
 import os
 from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple, Type
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Type
 from urllib.parse import urlparse
 
 from funcy import cached_property, collecting, project
@@ -34,6 +34,7 @@ from .utils import relpath
 from .utils.fs import path_isin
 
 if TYPE_CHECKING:
+    from dvc_data.hashfile.obj import HashFile
     from dvc_data.index import DataIndexKey
     from dvc_objects.db import ObjectDB
 
@@ -295,7 +296,7 @@ class Output:
         remote=None,
         repo=None,
         fs_config=None,
-        files=None,
+        files: List[Dict[str, Any]] = None,
     ):
         self.annot = Annotation(
             desc=desc, type=type, labels=labels or [], meta=meta or {}
@@ -344,7 +345,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._files = files
         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
@@ -353,7 +354,7 @@ class Output:
         self.live = live
 
         self.fs_path = self._parse_path(self.fs, fs_path)
-        self.obj = None
+        self.obj: Optional["HashFile"] = None
 
         self.remote = remote
 
@@ -374,10 +375,6 @@ class Output:
         )
         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)
@@ -423,6 +420,7 @@ class Output:
         self.hash_info = HashInfo.from_dict({})
         self.meta = Meta.from_dict({})
         self.obj = None
+        self._files = None
 
     @property
     def protocol(self):
@@ -622,6 +620,7 @@ class Output:
                 dry_run=True,
             )
             self.hash_info = obj.hash_info
+            self._files = None
             if not self.IS_DEPENDENCY:
                 logger.debug(
                     "Output '%s' doesn't use cache. Skipping saving.", self
@@ -638,6 +637,7 @@ class Output:
             ignore=self.dvcignore,
         )
         self.hash_info = self.obj.hash_info
+        self._files = None
 
     def set_exec(self):
         if self.isfile() and self.meta.isexec:
@@ -769,10 +769,13 @@ class Output:
             self.use_cache
             and self.is_in_repo
             and self.hash_info.isdir
-            and (kwargs.get("with_files") or self._dump_files)
+            and (kwargs.get("with_files") or self._files is not None)
         ):
-            assert self.obj
-            ret[self.PARAM_FILES] = self.obj.as_list(with_meta=True)
+            if self.obj:
+                obj = self.obj
+            else:
+                obj = self.get_obj()
+            ret[self.PARAM_FILES] = obj.as_list(with_meta=True)
 
         return ret
 
@@ -806,14 +809,22 @@ class Output:
         ) as cb:
             self.fs.get(self.fs_path, to.fs_path, batch_size=jobs, callback=cb)
 
-    def get_obj(self, filter_info=None, **kwargs):
+    def get_obj(
+        self, filter_info: Optional[str] = None, **kwargs
+    ) -> Optional["HashFile"]:
+        obj: Optional["HashFile"] = None
         if self.obj:
             obj = self.obj
         elif self.hash_info:
-            try:
-                obj = oload(self.odb, self.hash_info)
-            except (FileNotFoundError, ObjectFormatError):
-                return None
+            if self._files:
+                tree = Tree.from_list(self._files, hash_name=self.hash_name)
+                tree.digest()
+                obj = tree
+            else:
+                try:
+                    obj = oload(self.odb, self.hash_info)
+                except (FileNotFoundError, ObjectFormatError):
+                    return None
         else:
             return None
 
@@ -933,6 +944,7 @@ class Output:
         )
 
         self.hash_info = obj.hash_info
+        self._files = None
         return obj
 
     def get_files_number(self, filter_info=None):
@@ -1003,6 +1015,7 @@ class Output:
 
         obj = self.get_obj()
         if filter_info and filter_info != self.fs_path:
+            assert obj
             prefix = self.fs.path.parts(
                 self.fs.path.relpath(filter_info, self.fs_path)
             )
@@ -1139,6 +1152,7 @@ class Output:
         self.odb.add(merged.path, merged.fs, merged.oid)
 
         self.hash_info = merged.hash_info
+        self._files = None
         self.meta = Meta(
             size=du(self.odb, merged),
             nfiles=len(merged),