Add management that allows creating an archive from the contents of a…

… folder (#23)

cabinet/management/commands/archive_cabinet_folder.py

@@ -0,0 +1,48 @@
+import random
+import string
+from pathlib import Path
+from zipfile import ZIP_DEFLATED, ZipFile
+
+from cabinet.models import Folder
+from django.core.management import BaseCommand
+
+
+def _get_random_suffix():
+    return random.choices(string.ascii_lowercase, k=4)
+
+
+class Command(BaseCommand):
+    help = "Create archive with contents of a cabinet folder, using the data structure of the db instead of the disk."
+
+    def add_arguments(self, parser):
+        parser.add_argument("--folder-id", type=int, required=True)
+        parser.add_argument("--output", type=Path, required=True)
+
+    def handle(self, **options):
+        folder = Folder.objects.get(id=options["folder_id"])
+        output = options["output"]
+
+        arc_paths = set()
+        with ZipFile(output, "w", ZIP_DEFLATED) as zip_file:
+            for file, path in self._walk(folder, path=()):
+                arc_path = Path(*path) / file.file_name
+
+                if arc_path in arc_paths:
+                    filename = Path(file.file_name)
+                    arc_path = Path(*path) / "".join([
+                        filename.stem,
+                        "_",
+                        *_get_random_suffix(),
+                        *filename.suffixes,
+                    ])
+
+                zip_file.write(file.file.path, arc_path)
+                arc_paths.add(arc_path)
+
+    def _walk(self, folder, *, path):
+        path = (*tuple(path), folder.name)
+        for file_ in folder.files.all():
+            yield (file_, path)
+
+        for child in folder.children.all():
+            yield from self._walk(child, path=path)

tests/testapp/test_cabinet.py

@@ -1,3 +1,4 @@
+import tempfile
 import io
 import itertools
 import json
@@ -9,13 +10,17 @@
 from django.contrib.auth.models import User
 from django.core.exceptions import ImproperlyConfigured, ValidationError
 from django.core.files.base import ContentFile
+from django.core.management import call_command
 from django.test import Client, TestCase
 from django.test.utils import override_settings
 from django.urls import reverse
 
 from cabinet.base import AbstractFile, DownloadMixin, determine_accept_file_functions
 from cabinet.models import File, Folder, get_file_model
+from pathlib import Path
 from testapp.models import Stuff
+from unittest.mock import patch
+from zipfile import ZipFile
 
 
 class CabinetTestCase(TestCase):
@@ -547,3 +552,28 @@ class CustomFile(AbstractFile, NonModelMixin, DownloadMixin):
                 }
             ],
         )
+
+    @patch("cabinet.management.commands.archive_cabinet_folder._get_random_suffix", return_value="asdf")
+    def test_archive_management_command(self, patched__get_random_suffix):
+        output = "archive.zip"
+        folder = Folder.objects.create(name="Top")
+        subfolder = Folder.objects.create(parent=folder, name="Sub")
+
+        for _ in range(2):
+            file = File(folder=subfolder)
+            file.download_file.save("hello.txt", ContentFile("Hello"))
+
+        # enforce duplicate names
+        File.objects.all().update(file_name="hello.txt")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            output = Path(tmp_dir) / 'output.zip'
+            call_command('archive_cabinet_folder', folder_id=folder.id, output=output)
+            with ZipFile(output, "r") as zip_file:
+                self.assertEqual(
+                    zip_file.namelist(),
+                    [
+                        "Top/Sub/hello.txt",
+                        "Top/Sub/hello_asdf.txt",
+                    ]
+                )