feat(archive-output): improved study output un-archiving

antarest/study/storage/abstract_storage_service.py

@@ -3,6 +3,7 @@
 import shutil
 import tempfile
 import typing as t
+import zipfile
 from abc import ABC
 from pathlib import Path
 from uuid import uuid4
@@ -13,7 +14,7 @@
 from antarest.core.exceptions import BadOutputError, StudyOutputNotFoundError
 from antarest.core.interfaces.cache import CacheConstants, ICache
 from antarest.core.model import JSON, PublicMode
-from antarest.core.utils.utils import StopWatch, extract_archive, unzip, zip_dir
+from antarest.core.utils.utils import StopWatch, extract_archive, zip_dir
 from antarest.login.model import GroupDTO
 from antarest.study.common.studystorage import IStudyStorageService, T
 from antarest.study.model import (
@@ -251,9 +252,9 @@ def import_output(
                 elif output.suffix in {".zip", ".7z"}:
                     is_archived = True
                     path_output.rmdir()
-                    path_output = Path(str(path_output) + output.suffix)
+                    path_output = path_output.with_suffix(output.suffix)
                     shutil.copyfile(output, path_output)
-                    extension = ".zip" if output.suffix == ".zip" else ".7z"
+                    extension = output.suffix
             else:
                 extract_archive(output, path_output)
 
@@ -272,7 +273,7 @@ def import_output(
             logger.error("Failed to import output", exc_info=e)
             shutil.rmtree(path_output, ignore_errors=True)
             if is_archived:
-                Path(str(path_output) + ".zip").unlink(missing_ok=True)
+                path_output.with_suffix(extension).unlink(missing_ok=True)
             output_full_name = None
 
         return output_full_name
@@ -359,19 +360,20 @@ def unarchive_study_output(self, study: T, output_id: str, keep_src_archive: boo
                 f"Failed to archive study {study.name} output {output_id}. Maybe it's already unarchived",
             )
             return False
+
         try:
-            # use 7zip to uncompress the output folder
             if archive_path.suffix == ".7z":
-                with py7zr.SevenZipFile(archive_path, "r") as szf:
+                with py7zr.SevenZipFile(archive_path, mode="r") as szf:
                     szf.extractall(Path(study.path) / "output" / output_id)
-                if not keep_src_archive:
-                    archive_path.unlink()
+            elif archive_path.suffix == ".zip":
+                with zipfile.ZipFile(archive_path, mode="r") as zipf:
+                    zipf.extractall(Path(study.path) / "output" / output_id)
             else:
-                unzip(
-                    Path(study.path) / "output" / output_id,
-                    Path(study.path) / "output" / f"{output_id}.zip",
-                    remove_source_zip=not keep_src_archive,
-                )
+                raise NotImplementedError(f"Unsupported archive format {archive_path.suffix}")
+
+            if not keep_src_archive:
+                archive_path.unlink()
+
             remove_from_cache(self.cache, study.id)
             return True
         except Exception as e:

antarest/study/storage/utils.py

@@ -107,21 +107,21 @@ def find_single_output_path(all_output_path: Path) -> Path:
 def extract_output_name(path_output: Path, new_suffix_name: t.Optional[str] = None) -> str:
     ini_reader = IniReader()
     is_output_archived = path_output.suffix in {".zip", ".7z"}
+    info_filename = "info.antares-output"
     if is_output_archived:
-        temp_dir = tempfile.TemporaryDirectory()
-        s = StopWatch()
-        if path_output.suffix == ".zip":
-            with ZipFile(path_output, "r") as zip_obj:
-                zip_obj.extract("info.antares-output", temp_dir.name)
-        else:
-            with py7zr.SevenZipFile(path_output, mode="r") as archive:
-                archive.extract(targets=["info.antares-output"], path=temp_dir.name)
-        info_antares_output = ini_reader.read(Path(temp_dir.name) / "info.antares-output")
-        s.log_elapsed(lambda x: logger.info(f"info.antares_output has been read in {x}s"))
-        temp_dir.cleanup()
+        with tempfile.TemporaryDirectory() as temp_dir:
+            s = StopWatch()
+            if path_output.suffix == ".zip":
+                with ZipFile(path_output, mode="r") as zip_obj:
+                    zip_obj.extract(info_filename, temp_dir)
+            else:
+                with py7zr.SevenZipFile(path_output, mode="r") as archive:
+                    archive.extract(targets=[info_filename], path=temp_dir)
+            info_antares_output = ini_reader.read(Path(temp_dir) / info_filename)
+            s.log_elapsed(lambda x: logger.info(f"'{info_filename}' has been read in {x}s"))
 
     else:
-        info_antares_output = ini_reader.read(path_output / "info.antares-output")
+        info_antares_output = ini_reader.read(path_output / info_filename)
 
     general_info = info_antares_output["general"]
 
@@ -134,7 +134,7 @@ def extract_output_name(path_output: Path, new_suffix_name: t.Optional[str] = No
         general_info["name"] = suffix_name
         if not is_output_archived:
             ini_writer = IniWriter()
-            ini_writer.write(info_antares_output, path_output / "info.antares-output")
+            ini_writer.write(info_antares_output, path_output / info_filename)
         else:
             logger.warning("Could not rewrite the new name inside the output: the output is archived")