arkindex-base-worker 0.5.1rc1__tar.gz → 0.5.1rc2__tar.gz

{arkindex_base_worker-0.5.1rc1 → arkindex_base_worker-0.5.1rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arkindex-base-worker
-Version: 0.5.1rc1
+Version: 0.5.1rc2
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>

{arkindex_base_worker-0.5.1rc1 → arkindex_base_worker-0.5.1rc2}/arkindex_base_worker.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arkindex-base-worker
-Version: 0.5.1rc1
+Version: 0.5.1rc2
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>

{arkindex_base_worker-0.5.1rc1 → arkindex_base_worker-0.5.1rc2}/arkindex_worker/utils.py

@@ -4,6 +4,7 @@ import logging
 import os
 import tarfile
 import tempfile
+import zipfile
 from collections.abc import Callable, Generator
 from itertools import islice
 from pathlib import Path
@@ -225,6 +226,32 @@ def create_tar_zst_archive(
     return zst_fd, zst_archive, zst_hash, tar_hash
 
 
+def create_zip_archive(source: Path, destination: Path | None = None) -> Path:
+    """Helper to create a ZIP archive from a source folder.
+
+    :param source: Path to the folder whose content should be archived.
+    :param destination: Path to the created archive, defaults to None. If unspecified, a temporary file will be created.
+    :return: The file descriptor of the created tempfile (if one was created), path to the archive.
+    """
+    # Parse destination and create a tmpfile if none was specified
+    file_d, destination = (
+        tempfile.mkstemp(prefix="teklia-", suffix=".zip")
+        if destination is None
+        else (None, destination)
+    )
+    destination = Path(destination)
+    logger.debug(f"Compressing file to {destination}")
+
+    with zipfile.ZipFile(
+        destination, mode="w", compression=zipfile.ZIP_BZIP2
+    ) as archive:
+        for p in source.rglob("*"):
+            relpath = p.relative_to(source)
+            archive.write(p, arcname=relpath)
+
+        return archive, destination
+
+
 DEFAULT_BATCH_SIZE = 50
 """Batch size used for bulk publication to Arkindex"""
 

{arkindex_base_worker-0.5.1rc1 → arkindex_base_worker-0.5.1rc2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "arkindex-base-worker"
-version = "0.5.1rc1"
+version = "0.5.1rc2"
 description = "Base Worker to easily build Arkindex ML workflows"
 license = { file = "LICENSE" }
 dependencies = [

{arkindex_base_worker-0.5.1rc1 → arkindex_base_worker-0.5.1rc2}/tests/test_utils.py

@@ -1,4 +1,5 @@
 import logging
+import zipfile
 
 import pytest
 
@@ -7,6 +8,7 @@ from arkindex_worker.utils import (
     DEFAULT_BATCH_SIZE,
     batch_publication,
     close_delete_file,
+    create_zip_archive,
     extract_tar_zst_archive,
     parse_source_id,
 )
@@ -118,3 +120,43 @@ def test_batch_publication_decorator_alongside_unsupported_cache(caplog):
             "This API helper `custom_publication_in_batches_without_cache` did not update the cache database",
         ),
     ]
+
+
+def test_zip_archive():
+    # Create zip archive from fixtures
+    _, archive = create_zip_archive(FIXTURES_DIR / "extract_parent_archives/expected")
+
+    # Check the files in the archive
+    with zipfile.ZipFile(archive, mode="r") as f:
+        assert sorted(f.namelist()) == [
+            "test/",
+            "test/images/",
+            "test/images/f2649ce7-333e-44d2-ae73-387f18aad1f6.png",
+            "test/labels/",
+            "test/labels/f2649ce7-333e-44d2-ae73-387f18aad1f6.png",
+            "test/labels_json/",
+            "test/labels_json/f2649ce7-333e-44d2-ae73-387f18aad1f6.json",
+            "train/",
+            "train/images/",
+            "train/images/98115546-df07-448c-a2f0-34aa24789b77.png",
+            "train/images/ebeaa451-9287-4df7-9c40-07eb25cadb78.png",
+            "train/labels/",
+            "train/labels/98115546-df07-448c-a2f0-34aa24789b77.png",
+            "train/labels/ebeaa451-9287-4df7-9c40-07eb25cadb78.png",
+            "train/labels_json/",
+            "train/labels_json/98115546-df07-448c-a2f0-34aa24789b77.json",
+            "train/labels_json/ebeaa451-9287-4df7-9c40-07eb25cadb78.json",
+            "val/",
+            "val/images/",
+            "val/images/2987176d-4338-40f2-90d9-6d2cb4fd4a00.png",
+            "val/images/e3f91312-9201-45b7-9c32-e04a97ff1334.png",
+            "val/labels/",
+            "val/labels/2987176d-4338-40f2-90d9-6d2cb4fd4a00.png",
+            "val/labels/e3f91312-9201-45b7-9c32-e04a97ff1334.png",
+            "val/labels_json/",
+            "val/labels_json/2987176d-4338-40f2-90d9-6d2cb4fd4a00.json",
+            "val/labels_json/e3f91312-9201-45b7-9c32-e04a97ff1334.json",
+        ]
+
+    # Cleanup
+    archive.unlink()