diff --git a/antarest/study/service.py b/antarest/study/service.py index 119031d84e..7766a1e6dc 100644 --- a/antarest/study/service.py +++ b/antarest/study/service.py @@ -1,10 +1,10 @@ import base64 import contextlib +import fnmatch import io import json import logging import os -import re import time import typing as t from datetime import datetime, timedelta @@ -116,7 +116,7 @@ from antarest.study.storage.utils import ( MatrixProfile, assert_permission, - get_specific_matrices_according_to_version, + get_matrix_profile_by_version, get_start_date, is_managed, remove_from_cache, @@ -2426,26 +2426,25 @@ def get_matrix_with_index_and_header( ) -> pd.DataFrame: matrix_path = Path(path) study = self.get_study(study_id) - for aggregate in ["allocation", "correlation"]: - if matrix_path == Path("input") / "hydro" / aggregate: - all_areas = t.cast( - t.List[AreaInfoDTO], - self.get_all_areas(study_id, area_type=AreaType.AREA, ui=False, params=parameters), - ) - if aggregate == "allocation": - hydro_matrix = self.allocation_manager.get_allocation_matrix(study, all_areas) - else: - hydro_matrix = self.correlation_manager.get_correlation_matrix(all_areas, study, []) # type: ignore - return pd.DataFrame(data=hydro_matrix.data, columns=hydro_matrix.columns, index=hydro_matrix.index) - - json_matrix = self.get(study_id, path, depth=3, formatted=True, params=parameters) - for key in ["data", "index", "columns"]: - if key not in json_matrix: - raise IncorrectPathError(f"The path filled does not correspond to a matrix : {path}") - if not json_matrix["data"]: + + if matrix_path.parts in [("input", "hydro", "allocation"), ("input", "hydro", "correlation")]: + all_areas = t.cast( + t.List[AreaInfoDTO], + self.get_all_areas(study_id, area_type=AreaType.AREA, ui=False, params=parameters), + ) + if matrix_path.parts[-1] == "allocation": + hydro_matrix = self.allocation_manager.get_allocation_matrix(study, all_areas) + else: + hydro_matrix = self.correlation_manager.get_correlation_matrix(all_areas, study, []) # type: ignore + return pd.DataFrame(data=hydro_matrix.data, columns=hydro_matrix.columns, index=hydro_matrix.index) + + matrix_obj = self.get(study_id, path, depth=3, formatted=True, params=parameters) + if set(matrix_obj) != {"data", "index", "columns"}: + raise IncorrectPathError(f"The provided path does not point to a valid matrix: '{path}'") + if not matrix_obj["data"]: return pd.DataFrame() - df_matrix = pd.DataFrame(data=json_matrix["data"], columns=json_matrix["columns"], index=json_matrix["index"]) + df_matrix = pd.DataFrame(**matrix_obj) if with_index: matrix_index = self.get_input_matrix_startdate(study_id, path, parameters) time_column = pd.date_range( @@ -2453,14 +2452,15 @@ def get_matrix_with_index_and_header( ) df_matrix.index = time_column - specific_matrices = get_specific_matrices_according_to_version(int(study.version)) - for specific_matrix in specific_matrices: - if re.match(specific_matrix, path): + matrix_profiles = get_matrix_profile_by_version(int(study.version)) + for pattern, matrix_profile in matrix_profiles.items(): + if fnmatch.fnmatch(path, pattern): return _handle_specific_matrices( df_matrix, - specific_matrices[specific_matrix], + matrix_profile, path, with_index=with_index, with_header=with_header, ) + return df_matrix diff --git a/antarest/study/storage/utils.py b/antarest/study/storage/utils.py index db6d0bfe88..f3eca35ad1 100644 --- a/antarest/study/storage/utils.py +++ b/antarest/study/storage/utils.py @@ -239,30 +239,9 @@ def assert_permission( MATRIX_INPUT_DAYS_COUNT = 365 -MONTHS = ( - "January", - "February", - "March", - "April", - "May", - "June", - "July", - "August", - "September", - "October", - "November", - "December", -) +MONTHS = calendar.month_name[1:] -DAY_NAMES = ( - "Monday", - "Tuesday", - "Wednesday", - "Thursday", - "Friday", - "Saturday", - "Sunday", -) +DAY_NAMES = calendar.day_name[:] def _generate_columns(column_suffix: str) -> t.List[str]: @@ -378,12 +357,13 @@ class MatrixProfile(t.NamedTuple): SPECIFIC_MATRICES_870["input/bindingconstraints/*"] = MatrixProfile(cols=[], rows=[], stats=False) -def get_specific_matrices_according_to_version(study_version: int) -> t.Dict[str, MatrixProfile]: +def get_matrix_profile_by_version(study_version: int) -> t.Dict[str, MatrixProfile]: if study_version < 820: return SPECIFIC_MATRICES elif study_version < 870: return SPECIFIC_MATRICES_820 - return SPECIFIC_MATRICES_870 + else: + return SPECIFIC_MATRICES_870 def get_start_date( diff --git a/antarest/study/web/raw_studies_blueprint.py b/antarest/study/web/raw_studies_blueprint.py index 9a401d7135..d452a53e9e 100644 --- a/antarest/study/web/raw_studies_blueprint.py +++ b/antarest/study/web/raw_studies_blueprint.py @@ -317,7 +317,11 @@ def get_matrix( ) -> FileResponse: parameters = RequestParameters(user=current_user) df_matrix = study_service.get_matrix_with_index_and_header( - study_id=uuid, path=matrix_path, with_index=with_index, with_header=with_header, parameters=parameters + study_id=uuid, + path=matrix_path, + with_index=with_index, + with_header=with_header, + parameters=parameters, ) matrix_name = Path(matrix_path).stem diff --git a/tests/integration/raw_studies_blueprint/test_download_matrices.py b/tests/integration/raw_studies_blueprint/test_download_matrices.py index 9b9a269702..14b8aa9a58 100644 --- a/tests/integration/raw_studies_blueprint/test_download_matrices.py +++ b/tests/integration/raw_studies_blueprint/test_download_matrices.py @@ -209,7 +209,11 @@ def test_download_matrices(self, client: TestClient, admin_access_token: str, st assert dataframe.index[0] == "2018-01-01 00:00:00" dataframe.index = range(len(dataframe)) transposed_matrix = list(zip(*[8760 * [1.0], 8760 * [1.0], 8760 * [1.0], 8760 * [0.0]])) - expected_df = pd.DataFrame(columns=["0", "1", "2", "3"], index=range(8760), data=transposed_matrix) + expected_df = pd.DataFrame( + columns=["Marginal cost modulation", "Market bid modulation", "Capacity modulation", "Min gen modulation"], + index=range(8760), + data=transposed_matrix, + ) assert dataframe.equals(expected_df) # asserts endpoint returns the right columns for output matrix @@ -281,7 +285,7 @@ def test_download_matrices(self, client: TestClient, admin_access_token: str, st ) assert res.status_code == 404 assert res.json()["exception"] == "IncorrectPathError" - assert res.json()["description"] == "The path filled does not correspond to a matrix : settings/generaldata" + assert res.json()["description"] == "The provided path does not point to a valid matrix: 'settings/generaldata'" # wrong format res = client.get(