diff --git a/antarest/study/business/link_management.py b/antarest/study/business/link_management.py index a6ab65ca40..137120121a 100644 --- a/antarest/study/business/link_management.py +++ b/antarest/study/business/link_management.py @@ -20,10 +20,8 @@ from antarest.study.business.utils import execute_or_add_commands from antarest.study.model import RawStudy from antarest.study.storage.rawstudy.model.filesystem.config.links import ( - AssetType, LinkProperties, - LinkStyle, - TransmissionCapacity, + LinkStyle, TransmissionCapacity, AssetType, ) from antarest.study.storage.storage_service import StudyStorageService from antarest.study.storage.variantstudy.model.command.common import FilteringOptions @@ -41,14 +39,14 @@ class LinkInfoDTOBase(BaseModel): hurdles_cost: t.Optional[bool] = False loop_flow: t.Optional[bool] = False use_phase_shifter: t.Optional[bool] = False - transmission_capacities: t.Optional[str] = TransmissionCapacity.ENABLED - asset_type: t.Optional[str] = AssetType.AC + transmission_capacities: t.Optional[str] = TransmissionCapacity.ENABLED.value + asset_type: t.Optional[str] = AssetType.AC.value display_comments: t.Optional[bool] = True colorr: t.Optional[str] = DEFAULT_COLOR colorb: t.Optional[str] = DEFAULT_COLOR colorg: t.Optional[str] = DEFAULT_COLOR link_width: t.Optional[float] = 1 - link_style: t.Optional[str] = LinkStyle.PLAIN + link_style: t.Optional[str] = LinkStyle.PLAIN.value class LinkInfoDTO820(LinkInfoDTOBase): @@ -61,7 +59,7 @@ class LinkInfoDTO820(LinkInfoDTOBase): class LinkInfoFactory: @staticmethod - def create_link_info(version: int, **kwargs) -> LinkInfoDTOType: + def create_link_info(version: int, **kwargs: t.Any) -> LinkInfoDTOType: """ Creates a LinkInfoDTO object corresponding to the specified version. @@ -79,7 +77,7 @@ def create_link_info(version: int, **kwargs) -> LinkInfoDTOType: return link_info @staticmethod - def _initialize_link_info(version: int, **kwargs) -> LinkInfoDTOType: + def _initialize_link_info(version: int, **kwargs: t.Any) -> LinkInfoDTOType: """ Initializes the LinkInfoDTO object based on the study version. @@ -110,7 +108,7 @@ def _set_default_filters(version: int, link_info: LinkInfoDTOType) -> None: link_info.filter_year_by_year = FilteringOptions.FILTER_YEAR_BY_YEAR @staticmethod - def _check_version_coherence(version: int, **kwargs) -> None: + def _check_version_coherence(version: int, **kwargs: t.Any) -> None: """ Checks if filters are provided for a study version lower than 820. @@ -130,7 +128,7 @@ def _check_version_coherence(version: int, **kwargs) -> None: @staticmethod def create_parameters( study_version: int, link_creation_info: LinkInfoDTOType - ) -> t.Dict[str, t.Union[str, bool, float]]: + ) -> t.Dict[str, t.Union[str, bool, float, None]]: """ Creates the parameters for the link creation command, handling version differences. @@ -139,7 +137,7 @@ def create_parameters( link_creation_info (LinkInfoDTOType): The link information for creation. Returns: - t.Dict[str, t.Union[str, bool, float]: A dictionary containing the parameters for the command. + t.Dict[str, t.Union[str, bool, float, None]: A dictionary containing the parameters for the command. """ parameters = { "hurdles-cost": link_creation_info.hurdles_cost, @@ -217,6 +215,14 @@ def get_all_links(self, study: RawStudy, with_ui: bool = False) -> t.List[LinkIn def create_link(self, study: RawStudy, link_creation_info: LinkInfoDTOType) -> LinkInfoDTOType: study_version = int(study.version) + + if isinstance(link_creation_info, LinkInfoDTO820): + filter_synthesis = link_creation_info.filter_synthesis if study_version >= 820 else None + filter_year_by_year = link_creation_info.filter_year_by_year if study_version >= 820 else None + else: + filter_synthesis = None + filter_year_by_year = None + link_info_dto = LinkInfoFactory.create_link_info( version=int(study.version), area1=link_creation_info.area1, @@ -232,8 +238,8 @@ def create_link(self, study: RawStudy, link_creation_info: LinkInfoDTOType) -> L colorg=link_creation_info.colorg, link_width=link_creation_info.link_width, link_style=link_creation_info.link_style, - filter_synthesis=link_creation_info.filter_synthesis if study_version >= 820 else None, - filter_year_by_year=link_creation_info.filter_year_by_year if study_version >= 820 else None, + filter_synthesis=filter_synthesis, + filter_year_by_year=filter_year_by_year, ) storage_service = self.storage_service.get_storage(study) @@ -242,7 +248,7 @@ def create_link(self, study: RawStudy, link_creation_info: LinkInfoDTOType) -> L command = CreateLink( area1=link_creation_info.area1, area2=link_creation_info.area2, - parameters=LinkInfoFactory.create_parameters(int(study.version), link_info_dto), + parameters=LinkInfoFactory.create_parameters(study_version, link_info_dto), command_context=self.storage_service.variant_study_service.command_factory.command_context, ) diff --git a/antarest/study/storage/variantstudy/model/command/create_link.py b/antarest/study/storage/variantstudy/model/command/create_link.py index b1961efd5e..aea162b952 100644 --- a/antarest/study/storage/variantstudy/model/command/create_link.py +++ b/antarest/study/storage/variantstudy/model/command/create_link.py @@ -17,7 +17,7 @@ from antarest.core.model import JSON from antarest.core.utils.utils import assert_this from antarest.matrixstore.model import MatrixData -from antarest.study.storage.rawstudy.model.filesystem.config.links import AssetType, LinkStyle, TransmissionCapacity +from antarest.study.storage.rawstudy.model.filesystem.config.links import AssetType, TransmissionCapacity, LinkStyle from antarest.study.storage.rawstudy.model.filesystem.config.model import FileStudyTreeConfig, Link from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy from antarest.study.storage.variantstudy.business.utils import strip_matrix_protocol, validate_matrix @@ -31,9 +31,9 @@ class LinkProperties: LOOP_FLOW: bool = False USE_PHASE_SHIFTER: bool = False DISPLAY_COMMENTS: bool = True - TRANSMISSION_CAPACITIES: TransmissionCapacity = "enabled" - ASSET_TYPE: AssetType = "ac" - LINK_STYLE: str = "plain" + TRANSMISSION_CAPACITIES: str = TransmissionCapacity.ENABLED.value + ASSET_TYPE: str = AssetType.AC.value + LINK_STYLE: str = LinkStyle.PLAIN.value LINK_WIDTH: int = 1 COLORR: int = 112 COLORG: int = 112 diff --git a/tests/storage/business/test_arealink_manager.py b/tests/storage/business/test_arealink_manager.py index ac2ac275c5..e7bf272220 100644 --- a/tests/storage/business/test_arealink_manager.py +++ b/tests/storage/business/test_arealink_manager.py @@ -27,7 +27,6 @@ from antarest.study.business.link_management import ( LinkInfoDTO820, LinkInfoDTOBase, - LinkInfoDTOType, LinkInfoFactory, LinkManager, ) @@ -48,7 +47,6 @@ from antarest.study.storage.variantstudy.model.dbmodel import VariantStudy from antarest.study.storage.variantstudy.model.model import CommandDTO from antarest.study.storage.variantstudy.variant_study_service import VariantStudyService -from tests.conftest_services import study_storage_service_fixture from tests.helpers import with_db_context from tests.storage.business.assets import ASSETS_DIR