Read-tree extract tree from object

ugit/base.py

@@ -1,4 +1,4 @@
-from pathlib import Path
+from pathlib import Path, PurePath
 
 from . import data
 
@@ -24,5 +24,36 @@ def write_tree(directory="."):
     return data.hash_object(tree.encode(), "tree")
 
 
+def _iter_tree_entries(oid):
+    if not oid:
+        return
+    tree = data.get_object(oid, "tree")
+    for entry in tree.decode().splitlines():
+        type_, oid, name = entry.split(" ", 2)
+        yield type_, oid, name
+
+
+def get_tree(oid, base_path=""):
+    result = {}
+    for type_, oid, name in _iter_tree_entries(oid):
+        assert "/" not in name
+        assert name not in ("..", ".")
+        path = base_path + name
+        if type_ == "blob":
+            result[path] = oid
+        elif type_ == "tree":
+            result.update(get_tree(oid, f"{path}/"))
+        else:
+            assert False, f"Unknown tree entry {type_}"
+    return result
+
+
+def read_tree(tree_oid):
+    for path, oid in get_tree(tree_oid, base_path="./").items():
+        Path.mkdir(PurePath.parent(path), exist_ok=True)
+        with open(path, "wb") as f:
+            f.write(data.get_object(oid))
+
+
 def is_ignored(path):
     return ".ugit" in path.split("/")

ugit/cli.py

@@ -21,6 +21,10 @@ def parse_args():
     init_parser = commands.add_parser("init")
     init_parser.set_defaults(func=init)
 
+    hash_object_parser = commands.add_parser("hash-object")
+    hash_object_parser.set_defaults(func=hash_object)
+    hash_object_parser.add_argument("file")
+
     cat_file_parser = commands.add_parser("cat-file")
     cat_file_parser.set_defaults(func=cat_file)
     cat_file_parser.add_argument("object")
@@ -28,9 +32,9 @@ def parse_args():
     write_tree_parser = commands.add_parser("write-tree")
     write_tree_parser.set_defaults(func=write_tree)
 
-    hash_object_parser = commands.add_parser("hash-object")
-    hash_object_parser.set_defaults(func=hash_object)
-    hash_object_parser.add_argument("file")
+    read_tree_parser = commands.add_parser("read-tree")
+    read_tree_parser.set_defaults(func=read_tree)
+    read_tree_parser.add_argument("tree")
 
     return parser.parse_args()
 
@@ -52,3 +56,7 @@ def cat_file(args):
 
 def write_tree(args):
     print(base.write_tree())
+
+
+def read_tree(args):
+    base.read_tree(args.tree)