objectstore: introduce TreeObject

Introduce a small `TreeObject` class that is the representation of
a tree during its construction. It supports calculating its treesum
as well initialize the new tree with an existing one.

osbuild/objectstore.py

@@ -26,6 +26,31 @@ def suppress_oserror(*errnos):
             raise e
 
 
+class TreeObject:
+    def __init__(self, path: str):
+        os.makedirs(path, mode=0o755, exist_ok=True)
+        self.path = path
+
+    def init(self, source: str) -> None:
+        """Initialize the tree with a source tree"""
+        subprocess.run(["cp", "--reflink=auto", "-a",
+                        f"{source}/.", self.path],
+                       check=True)
+
+    @property
+    def treesum(self) -> str:
+        """Calculate the treesum of the tree"""
+        fd = os.open(self.path, os.O_DIRECTORY)
+        try:
+            m = hashlib.sha256()
+            treesum.treesum(m, fd)
+            treesum_hash = m.hexdigest()
+        finally:
+            os.close(fd)
+
+        return treesum_hash
+
+
 class ObjectStore:
     def __init__(self, store):
         self.store = store
@@ -53,6 +78,7 @@ class ObjectStore:
                 # None was given as object_id, just return an empty directory
                 yield tmp
 
+
     @contextlib.contextmanager
     def new(self, object_id, base_id=None):
         """Creates a new directory for `object_id`.
@@ -64,28 +90,21 @@ class ObjectStore:
         with tempfile.TemporaryDirectory(dir=self.store) as tmp:
             # the tree that is yielded will be added to the content store
             # on success as object_id
-
-            tree = f"{tmp}/tree"
-            link = f"{tmp}/link"
-            os.mkdir(tree, mode=0o755)
+            tree = TreeObject(f"{tmp}/tree")
 
             if base_id:
                 # the base, the working tree and the output tree are all on
                 # the same fs, so attempt a lightweight copy if the fs
                 # supports it
-                subprocess.run(["cp", "--reflink=auto", "-a", f"{self.refs}/{base_id}/.", tree], check=True)
+                tree.init(f"{self.refs}/{base_id}")
 
-            yield tree
+            yield tree.path
+
+            # if the yield above raises an exception, the working tree
+            # is cleaned up by tempfile, otherwise, we save it in the
+            # correct place:
+            treesum_hash = tree.treesum
 
-            # if the yield raises an exception, the working tree is cleaned
-            # up by tempfile, otherwise, we save it in the correct place:
-            fd = os.open(tree, os.O_DIRECTORY)
-            try:
-                m = hashlib.sha256()
-                treesum.treesum(m, fd)
-                treesum_hash = m.hexdigest()
-            finally:
-                os.close(fd)
             # the tree is stored in the objects directory using its content
             # hash as its name, ideally a given object_id (i.e., given config)
             # will always produce the same content hash, but that is not
@@ -95,10 +114,12 @@ class ObjectStore:
             # if a tree with the same treesum already exist, use that
             with suppress_oserror(errno.ENOTEMPTY):
                 os.rename(tree, output_tree)
+            tree.path = output_tree
 
             # symlink the object_id (config hash) in the refs directory to the
             # treesum (content hash) in the objects directory. If a symlink by
             # that name alreday exists, atomically replace it, but leave the
             # backing object in place (it may be in use).
+            link = f"{tmp}/link"
             os.symlink(f"../objects/{treesum_hash}", link)
             os.replace(link, f"{self.refs}/{object_id}")