to_disk: Add type annotations + a simple test for DirectoryBuilder

swh/vault/tests/test_to_disk.py

-# Copyright (C) 2020  The Software Heritage developers
+# Copyright (C) 2020-2022  The Software Heritage developers
 # See the AUTHORS file at the top-level directory of this distribution
 # License: GNU General Public License version 3, or any later version
 # See top-level LICENSE file for more information
 
 import pytest
 
-from swh.model.model import Content, SkippedContent
-from swh.vault.to_disk import get_filtered_files_content
+from swh.model.from_disk import DentryPerms
+from swh.model.model import Content, Directory, DirectoryEntry, SkippedContent
+from swh.vault.to_disk import DirectoryBuilder, get_filtered_files_content
 
 
 def test_get_filtered_files_content(swh_storage):
@@ -76,3 +77,58 @@ def test_get_filtered_files_content__unknown_status(swh_storage):
 
     with pytest.raises(AssertionError, match="unexpected status None"):
         list(get_filtered_files_content(swh_storage, files_data))
+
+
+def test_directory_builder(swh_storage, tmp_path):
+    cnt1 = Content.from_data(b"foo bar")
+    cnt2 = Content.from_data(b"bar baz")
+    cnt3 = Content.from_data(b"baz qux")
+    dir1 = Directory(
+        entries=(
+            DirectoryEntry(
+                name=b"content1",
+                type="file",
+                target=cnt1.sha1_git,
+                perms=DentryPerms.content,
+            ),
+            DirectoryEntry(
+                name=b"content2",
+                type="file",
+                target=cnt2.sha1_git,
+                perms=DentryPerms.content,
+            ),
+        )
+    )
+    dir2 = Directory(
+        entries=(
+            DirectoryEntry(
+                name=b"content3",
+                type="file",
+                target=cnt3.sha1_git,
+                perms=DentryPerms.content,
+            ),
+            DirectoryEntry(
+                name=b"subdirectory",
+                type="dir",
+                target=dir1.id,
+                perms=DentryPerms.directory,
+            ),
+        )
+    )
+    swh_storage.content_add([cnt1, cnt2, cnt3])
+    swh_storage.directory_add([dir1, dir2])
+
+    root = tmp_path / "root"
+    builder = DirectoryBuilder(swh_storage, bytes(root), dir2.id)
+
+    assert not root.exists()
+
+    builder.build()
+
+    assert root.is_dir()
+    assert set(root.glob("**/*")) == {
+        root / "subdirectory",
+        root / "subdirectory" / "content1",
+        root / "subdirectory" / "content2",
+        root / "content3",
+    }

swh/vault/to_disk.py

@@ -71,7 +71,7 @@ def apply_chunked(func, input_list, chunk_size):
 class DirectoryBuilder:
     """Reconstructs the on-disk representation of a directory in the storage."""
 
-    def __init__(self, storage, root, dir_id):
+    def __init__(self, storage: StorageInterface, root: bytes, dir_id: bytes):
         """Initialize the directory builder.
 
         Args:
@@ -83,7 +83,7 @@ class DirectoryBuilder:
         self.root = root
         self.dir_id = dir_id
 
-    def build(self):
+    def build(self) -> None:
         """Perform the reconstruction of the directory in the given root."""
         # Retrieve data from the database.
         # Split into files, revisions and directory data.
@@ -96,7 +96,7 @@ class DirectoryBuilder:
         self._create_files(entries["file"])
         self._create_revisions(entries["rev"])
 
-    def _create_tree(self, directories):
+    def _create_tree(self, directories: List[Dict[str, Any]]) -> None:
         """Create a directory tree from the given paths
 
         The tree is created from `root` and each given directory in
@@ -109,7 +109,7 @@ class DirectoryBuilder:
         for dir in directories:
             os.makedirs(os.path.join(self.root, dir["path"]))
 
-    def _create_files(self, files_data):
+    def _create_files(self, files_data: List[Dict[str, Any]]) -> None:
         """Create the files in the tree and fetch their contents."""
         f = functools.partial(get_filtered_files_content, self.storage)
         files_data = apply_chunked(f, files_data, 1000)
@@ -118,7 +118,7 @@ class DirectoryBuilder:
             path = os.path.join(self.root, file_data["path"])
             self._create_file(path, file_data["content"], file_data["perms"])
 
-    def _create_revisions(self, revs_data):
+    def _create_revisions(self, revs_data: List[Dict[str, Any]]) -> None:
         """Create the revisions in the tree as broken symlinks to the target
         identifier."""
         for file_data in revs_data:
@@ -126,7 +126,9 @@ class DirectoryBuilder:
             target = hashutil.hash_to_hex(file_data["target"])
             self._create_file(path, target, mode=DentryPerms.symlink)
 
-    def _create_file(self, path, content, mode=DentryPerms.content):
+    def _create_file(
+        self, path: bytes, content: bytes, mode: int = DentryPerms.content
+    ) -> None:
         """Create the given file and fill it with content."""
         perms = mode_to_perms(mode)
         if perms == DentryPerms.symlink: