From 98f7d1e481afb84e7bca858d4c990e32fe9bd198 Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <batuhanosmantaskaya@gmail.com>
Date: Tue, 6 Jul 2021 10:56:13 +0300
Subject: [PATCH] objects: hash_file: retrieve size from hash_info (#6284)

* objects: hash_file: retrieve size from hash_info

* objects: drop Tree.size implementation
---
 dvc/objects/file.py             |  4 +---
 dvc/objects/tree.py             | 12 ++++--------
 tests/unit/objects/test_tree.py | 25 +++++++++++++++++--------
 3 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/dvc/objects/file.py b/dvc/objects/file.py
index f39ceed8c..4fcf6fd5f 100644
--- a/dvc/objects/file.py
+++ b/dvc/objects/file.py
@@ -28,9 +28,7 @@ class HashFile:
 
     @property
     def size(self):
-        if not (self.path_info and self.fs):
-            return None
-        return self.fs.getsize(self.path_info)
+        return self.hash_info.size
 
     def __len__(self):
         return 1
diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py
index 783a41366..0debbfc09 100644
--- a/dvc/objects/tree.py
+++ b/dvc/objects/tree.py
@@ -28,13 +28,6 @@ class Tree(HashFile):
 
         return Trie(self._dict)
 
-    @property
-    def size(self):
-        try:
-            return sum(obj.size for _, obj in self)
-        except TypeError:
-            return None
-
     def add(self, key, obj):
         self.__dict__.pop("trie", None)
         self._dict[key] = obj
@@ -52,7 +45,10 @@ class Tree(HashFile):
         self.path_info = path_info
         self.hash_info = get_file_hash(path_info, memfs, "md5")
         self.hash_info.value += ".dir"
-        self.hash_info.size = self.size
+        try:
+            self.hash_info.size = sum(obj.size for _, obj in self)
+        except TypeError:
+            self.hash_info.size = None
         self.hash_info.nfiles = len(self)
 
     def __len__(self):
diff --git a/tests/unit/objects/test_tree.py b/tests/unit/objects/test_tree.py
index 41c567e00..fc6953451 100644
--- a/tests/unit/objects/test_tree.py
+++ b/tests/unit/objects/test_tree.py
@@ -69,19 +69,27 @@ def test_list(lst, trie_dict):
         ({}, 0),
         (
             {
-                ("a",): HashInfo("md5", "abc", size=1),
-                ("b",): HashInfo("md5", "def", size=2),
-                ("c",): HashInfo("md5", "ghi", size=3),
-                ("dir", "foo"): HashInfo("md5", "jkl", size=4),
-                ("dir", "bar"): HashInfo("md5", "mno", size=5),
-                ("dir", "baz"): HashInfo("md5", "pqr", size=6),
+                ("a",): HashFile(None, None, HashInfo("md5", "abc", size=1)),
+                ("b",): HashFile(None, None, HashInfo("md5", "def", size=2)),
+                ("c",): HashFile(None, None, HashInfo("md5", "ghi", size=3)),
+                ("dir", "foo"): HashFile(
+                    None, None, HashInfo("md5", "jkl", size=4)
+                ),
+                ("dir", "bar"): HashFile(
+                    None, None, HashInfo("md5", "mno", size=5)
+                ),
+                ("dir", "baz"): HashFile(
+                    None, None, HashInfo("md5", "pqr", size=6)
+                ),
             },
             21,
         ),
         (
             {
-                ("a",): HashInfo("md5", "abc", size=1),
-                ("b",): HashInfo("md5", "def", size=None),
+                ("a",): HashFile(None, None, HashInfo("md5", "abc", size=1)),
+                ("b",): HashFile(
+                    None, None, HashInfo("md5", "def", size=None)
+                ),
             },
             None,
         ),
@@ -90,6 +98,7 @@ def test_list(lst, trie_dict):
 def test_size(trie_dict, size):
     tree = Tree(None, None, None)
     tree._dict = trie_dict
+    tree.digest()
     assert tree.size == size