From a5cd030227c60830b2ae621fb2f341aa64d885c4 Mon Sep 17 00:00:00 2001 From: Laurent LAPORTE Date: Wed, 5 Jun 2024 22:54:20 +0200 Subject: [PATCH] feat(archive-output): improved study output un-archiving --- .../study/storage/abstract_storage_service.py | 28 ++++++++++--------- antarest/study/storage/utils.py | 26 ++++++++--------- 2 files changed, 28 insertions(+), 26 deletions(-) diff --git a/antarest/study/storage/abstract_storage_service.py b/antarest/study/storage/abstract_storage_service.py index f276c0d227..65846df71f 100644 --- a/antarest/study/storage/abstract_storage_service.py +++ b/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: diff --git a/antarest/study/storage/utils.py b/antarest/study/storage/utils.py index 6464385b3c..fb7584f7a2 100644 --- a/antarest/study/storage/utils.py +++ b/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")