From b03d3b4a347edb757b3685cfe602d2cbdec1d14f Mon Sep 17 00:00:00 2001 From: belthlemar Date: Wed, 11 Oct 2023 16:01:14 +0200 Subject: [PATCH] feat(utils): add 7z support for matrices import --- antarest/matrixstore/service.py | 23 ++++++-- tests/integration/assets/matrices.7z | Bin 0 -> 394 bytes tests/integration/assets/matrices.zip | Bin 0 -> 2518 bytes tests/integration/test_integration.py | 79 ++++++++++++++++++-------- 4 files changed, 71 insertions(+), 31 deletions(-) create mode 100644 tests/integration/assets/matrices.7z create mode 100644 tests/integration/assets/matrices.zip diff --git a/antarest/matrixstore/service.py b/antarest/matrixstore/service.py index 639084b587..a86bd2cfce 100644 --- a/antarest/matrixstore/service.py +++ b/antarest/matrixstore/service.py @@ -12,6 +12,7 @@ import numpy as np from fastapi import UploadFile from numpy import typing as npt +from py7zr import SevenZipFile from antarest.core.config import Config from antarest.core.filetransfer.model import FileDownloadTaskDTO @@ -187,12 +188,22 @@ def create_by_importation(self, file: UploadFile, is_json: bool = False) -> List with contextlib.closing(f): buffer = io.BytesIO(f.read()) matrix_info: List[MatrixInfoDTO] = [] - with zipfile.ZipFile(buffer) as zf: - for info in zf.infolist(): - if info.is_dir() or info.filename in EXCLUDED_FILES: - continue - matrix_id = self._file_importation(zf.read(info.filename), is_json=is_json) - matrix_info.append(MatrixInfoDTO(id=matrix_id, name=info.filename)) + if file.filename.endswith("zip"): + with zipfile.ZipFile(buffer) as zf: + for info in zf.infolist(): + if info.is_dir() or info.filename in EXCLUDED_FILES: + continue + matrix_id = self._file_importation(zf.read(info.filename), is_json=is_json) + matrix_info.append(MatrixInfoDTO(id=matrix_id, name=info.filename)) + else: + with SevenZipFile(buffer, "r") as szf: + for info in szf.list(): + if info.is_directory or info.filename in EXCLUDED_FILES: # type:ignore + continue + file_content = next(iter(szf.read(info.filename).values())) + matrix_id = self._file_importation(file_content.read(), is_json=is_json) + matrix_info.append(MatrixInfoDTO(id=matrix_id, name=info.filename)) + szf.reset() return matrix_info else: matrix_id = self._file_importation(f.read(), is_json=is_json) diff --git a/tests/integration/assets/matrices.7z b/tests/integration/assets/matrices.7z new file mode 100644 index 0000000000000000000000000000000000000000..2d4a69a4d61ed33d37000744d3a8f45da7ea6741 GIT binary patch literal 394 zcmV;50d@X2dc3bE8~_8><{=J80RR910000Z000000001^2O+fp7>qE1GKv@sDNdyn zE;ON>>0oFx&kY+~@E^ymt$K#aen`Qi(+-x z!ExXsXn>)3qWgn)1i1w4`RWrU^hxsu6;hm4FcGL;=2eftL2mgTllXCXS@;Xh|CK04 z52^fo-us#`E>^&gh4X`fz-7cGQ~33fHZcXIu+c zW>Q$b>2|DHE^g{GCJp|8TGX0001iGY75@)iE8}mn75vvwkm$lD@t?+L&c& zE|tf5EgvBJ(5;!gKM`4DF+Auj5*Cc&FPG zO_th&enBHDBC!cc5gh<7uklF;4p)Vf%C3shxxpd?+oeGc&^-*icJe)l)dHC#VE^m| oH~<$0fWiR@fQ0}D3jqKDBLe{e1zi9DfB+1Dn+gH_#EoSD09^^PU;qFB literal 0 HcmV?d00001 diff --git a/tests/integration/assets/matrices.zip b/tests/integration/assets/matrices.zip new file mode 100644 index 0000000000000000000000000000000000000000..98ac501644f7894e2185be1cbd19c49b81505931 GIT binary patch literal 2518 zcmWIWW@Zs#-~d9$N$P9G%1Xt(GVC7 zfzc3vg~0xU%mLo)97dLxu34}!Fn~&&0K8>UW(i7Jq{m{Ia&N+DSp*Bu(E@ff1V%%E zyb!=z7BMo3Fe6$<$mJBMWds9T8bK`L+DQT4#G2UyED?z-?*hD8*+4E}0>TwQT7nbA F0|3-OLy`ah literal 0 HcmV?d00001 diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py index 376a8ffc80..b3aa088c82 100644 --- a/tests/integration/test_integration.py +++ b/tests/integration/test_integration.py @@ -224,31 +224,6 @@ def test_main(client: TestClient, admin_access_token: str, study_id: str) -> Non ) assert len(res.json()) == 4 - # tests outputs import for .zip - output_path_zip = ASSETS_DIR / "output_adq.zip" - client.post( - f"/v1/studies/{study_id}/output", - headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, - files={"output": io.BytesIO(output_path_zip.read_bytes())}, - ) - res = client.get( - f"/v1/studies/{study_id}/outputs", - headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, - ) - assert len(res.json()) == 5 - # tests outputs import for .7z - output_path_seven_zip = ASSETS_DIR / "output_adq.7z" - client.post( - f"/v1/studies/{study_id}/output", - headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, - files={"output": io.BytesIO(output_path_seven_zip.read_bytes())}, - ) - res = client.get( - f"/v1/studies/{study_id}/outputs", - headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, - ) - assert len(res.json()) == 6 - # study creation created = client.post( "/v1/studies?name=foo", @@ -2466,6 +2441,60 @@ def test_import(client: TestClient, admin_access_token: str, study_id: str) -> N ) assert res.status_code == 201 + # tests outputs import for .zip + output_path_zip = ASSETS_DIR / "output_adq.zip" + client.post( + f"/v1/studies/{study_id}/output", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + files={"output": io.BytesIO(output_path_zip.read_bytes())}, + ) + res = client.get( + f"/v1/studies/{study_id}/outputs", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + ) + assert len(res.json()) == 6 + + # tests outputs import for .7z + output_path_seven_zip = ASSETS_DIR / "output_adq.7z" + client.post( + f"/v1/studies/{study_id}/output", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + files={"output": io.BytesIO(output_path_seven_zip.read_bytes())}, + ) + res = client.get( + f"/v1/studies/{study_id}/outputs", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + ) + assert len(res.json()) == 7 + + # test matrices import for .zip file + matrices_zip_path = ASSETS_DIR / "matrices.zip" + res = client.post( + "/v1/matrix/_import", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + files={"file": (matrices_zip_path.name, io.BytesIO(matrices_zip_path.read_bytes()), "application/zip")}, + ) + assert res.status_code == 200 + result = res.json() + assert len(result) == 2 + assert result[0]["name"] == "fr.txt" + assert result[1]["name"] == "it.txt" + + # test matrices import for .7z file + matrices_seven_zip_path = ASSETS_DIR / "matrices.7z" + res = client.post( + "/v1/matrix/_import", + headers={"Authorization": f'Bearer {george_credentials["access_token"]}'}, + files={ + "file": (matrices_seven_zip_path.name, io.BytesIO(matrices_seven_zip_path.read_bytes()), "application/zip") + }, + ) + assert res.status_code == 200 + result = res.json() + assert len(result) == 2 + assert result[0]["name"] == "fr.txt" + assert result[1]["name"] == "it.txt" + def test_copy(client: TestClient, admin_access_token: str, study_id: str) -> None: admin_headers = {"Authorization": f"Bearer {admin_access_token}"}