fix(mypy): configure the use of pydantic plugin

otherwise options such as populate_by_name are not correctly taken into account Signed-off-by: Sylvain Leclerc <[email protected]>
AntaresSimulatorTeam · Sep 5, 2024 · b4ed725 · b4ed725
1 parent bc5a089
commit b4ed725
Show file tree

Hide file tree

Showing 15 changed files with 155 additions and 29 deletions.
diff --git a/antarest/launcher/model.py b/antarest/launcher/model.py
@@ -15,7 +15,7 @@
 import typing as t
 from datetime import datetime
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, Sequence, String  # type: ignore
 from sqlalchemy.orm import relationship  # type: ignore
 
@@ -122,7 +122,24 @@ class JobResultDTO(BaseModel):
     exit_code: t.Optional[int]
     solver_stats: t.Optional[str]
     owner: t.Optional[UserInfo]
-    # TODO SL: restore example ?
+
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = JobResultDTO(
+                id="b2a9f6a7-7f8f-4f7a-9a8b-1f9b4c5d6e7f",
+                study_id="b2a9f6a7-7f8f-4f7a-9a8b-1f9b4c5d6e7f",
+                launcher="slurm",
+                launcher_params='{"nb_cpu": 4, "time_limit": 3600}',
+                status=JobStatus.SUCCESS,
+                creation_date="2023-11-25 12:00:00",
+                completion_date="2023-11-25 12:27:31",
+                msg="Study successfully executed",
+                output_id="20231125-1227eco",
+                exit_code=0,
+                solver_stats="time: 1651s, call_count: 1, optimization_issues: []",
+                owner=UserInfo(id=0o007, name="James BOND"),
+            ).model_dump()
 
 
 class JobLog(Base):  # type: ignore

diff --git a/antarest/main.py b/antarest/main.py
@@ -286,6 +286,10 @@ def home(request: Request) -> Any:
         def home(request: Request) -> Any:
             return ""
 
+    # TODO SL:         on_event is deprecated, use lifespan event handlers instead.
+    #
+    #         Read more about it in the
+    #         [FastAPI docs for Lifespan Events](https://fastapi.tiangolo.com/advanced/events/).
     @application.on_event("startup")
     def set_default_executor() -> None:
         import asyncio

diff --git a/antarest/matrixstore/matrix_editor.py b/antarest/matrixstore/matrix_editor.py
@@ -158,7 +158,14 @@ class MatrixEditInstruction(BaseModel):
 
     class Config:
         extra = "forbid"
-        # TODO SL: restore schema extra ?
+        json_schema_extra = {
+            "example": {
+                "column_from": 5,
+                "column_to": 8,
+                "row_from": 0,
+                "row_to": 8760,
+            }
+        }
 
     @model_validator(mode="before")
     def check_slice_coordinates(cls, values: Dict[str, Any]) -> Dict[str, Any]:

diff --git a/antarest/matrixstore/service.py b/antarest/matrixstore/service.py
@@ -164,7 +164,7 @@ def _from_dto(dto: MatrixDTO) -> t.Tuple[Matrix, MatrixContent]:
             created_at=datetime.fromtimestamp(dto.created_at),
         )
 
-        content = MatrixContent(data=dto.data, index=dto.index, columns=dto.columns)  # type: ignore
+        content = MatrixContent(data=dto.data, index=dto.index, columns=dto.columns)
 
         return matrix, content
 

diff --git a/antarest/study/business/all_optional_meta.py b/antarest/study/business/all_optional_meta.py
@@ -10,8 +10,8 @@
 #
 # This file is part of the Antares project.
 
-import typing as t
 import copy
+import typing as t
 
 from pydantic import BaseModel, create_model
 
@@ -40,6 +40,10 @@ def camel_case_model(model: t.Type[BaseModel]) -> t.Type[BaseModel]:
         The modified model.
     """
     model.model_config["alias_generator"] = to_camel_case
+
+    # Manually overriding already defined alias names (in base classes),
+    # otherwise they have precedence over generated ones.
+    # TODO There is probably a better way to handle those cases
     for field_name, field in model.model_fields.items():
         new_alias = to_camel_case(field_name)
         field.alias = new_alias

diff --git a/antarest/study/business/area_management.py b/antarest/study/business/area_management.py
@@ -660,7 +660,7 @@ def update_area_metadata(
         self.patch_service.save(study, patch)
         return AreaInfoDTO(
             id=area_id,
-            name=area_or_set.name if area_or_set is not None else area_id,  # type: ignore
+            name=area_or_set.name if area_or_set is not None else area_id,
             type=AreaType.AREA if isinstance(area_or_set, Area) else AreaType.DISTRICT,
             metadata=patch.areas.get(area_id),
             set=area_or_set.get_areas(list(file_study.config.areas)) if isinstance(area_or_set, DistrictSet) else [],

diff --git a/antarest/study/business/areas/renewable_management.py b/antarest/study/business/areas/renewable_management.py
@@ -45,14 +45,27 @@ class TimeSeriesInterpretation(EnumIgnoreCase):
     PRODUCTION_FACTOR = "production-factor"
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class RenewableClusterInput(RenewableProperties):
     """
     Model representing the data structure required to edit an existing renewable cluster.
     """
 
+    class Config:
+        populate_by_name = True
+
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = RenewableClusterInput(
+                group="Gas",
+                name="Gas Cluster XY",
+                enabled=False,
+                unit_count=100,
+                nominal_capacity=1000.0,
+                ts_interpretation="power-generation",
+            ).model_dump()
+
 
 class RenewableClusterCreation(RenewableClusterInput):
     """
@@ -74,14 +87,26 @@ def to_config(self, study_version: t.Union[str, int]) -> RenewableConfigType:
         return create_renewable_config(study_version=study_version, **values)
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class RenewableClusterOutput(RenewableConfig):
     """
     Model representing the output data structure to display the details of a renewable cluster.
     """
 
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = RenewableClusterOutput(
+                id="Gas cluster YZ",
+                group="Gas",
+                name="Gas Cluster YZ",
+                enabled=False,
+                unit_count=100,
+                nominal_capacity=1000.0,
+                ts_interpretation="power-generation",
+            ).model_dump()
+
 
 def create_renewable_output(
     study_version: t.Union[str, int],

diff --git a/antarest/study/business/areas/st_storage_management.py b/antarest/study/business/areas/st_storage_management.py
@@ -36,6 +36,7 @@
     STStorage880Config,
     STStorage880Properties,
     STStorageConfigType,
+    STStorageGroup,
     create_st_storage_config,
 )
 from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
@@ -46,14 +47,27 @@
 from antarest.study.storage.variantstudy.model.command.update_config import UpdateConfig
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class STStorageInput(STStorage880Properties):
     """
     Model representing the form used to EDIT an existing short-term storage.
     """
 
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = STStorageInput(
+                name="Siemens Battery",
+                group=STStorageGroup.BATTERY,
+                injection_nominal_capacity=150,
+                withdrawal_nominal_capacity=150,
+                reservoir_capacity=600,
+                efficiency=0.94,
+                initial_level=0.5,
+                initial_level_optim=True,
+            ).model_dump()
+
 
 class STStorageCreation(STStorageInput):
     """
@@ -76,14 +90,27 @@ def to_config(self, study_version: t.Union[str, int]) -> STStorageConfigType:
         return create_st_storage_config(study_version=study_version, **values)
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class STStorageOutput(STStorage880Config):
     """
     Model representing the form used to display the details of a short-term storage entry.
     """
 
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = STStorageOutput(
+                id="siemens_battery",
+                name="Siemens Battery",
+                group=STStorageGroup.BATTERY,
+                injection_nominal_capacity=150,
+                withdrawal_nominal_capacity=150,
+                reservoir_capacity=600,
+                efficiency=0.94,
+                initial_level_optim=True,
+            ).model_dump()
+
 
 # =============
 #  Time series

diff --git a/antarest/study/business/areas/thermal_management.py b/antarest/study/business/areas/thermal_management.py
@@ -53,15 +53,28 @@
 _ALL_CLUSTERS_PATH = "input/thermal/clusters"
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class ThermalClusterInput(Thermal870Properties):
     """
     Model representing the data structure required to edit an existing thermal cluster within a study.
     """
 
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = ThermalClusterInput(
+                group="Gas",
+                name="Gas Cluster XY",
+                enabled=False,
+                unit_count=100,
+                nominal_capacity=1000.0,
+                gen_ts="use global",
+                co2=7.0,
+            ).model_dump()
 
+
+@camel_case_model
 class ThermalClusterCreation(ThermalClusterInput):
     """
     Model representing the data structure required to create a new thermal cluster within a study.
@@ -82,14 +95,27 @@ def to_config(self, study_version: t.Union[str, int]) -> ThermalConfigType:
         return create_thermal_config(study_version=study_version, **values)
 
 
-# TODO SL : restore schema ?
 @all_optional_model
 @camel_case_model
 class ThermalClusterOutput(Thermal870Config):
     """
     Model representing the output data structure to display the details of a thermal cluster within a study.
     """
 
+    class Config:
+        @staticmethod
+        def json_schema_extra(schema: t.MutableMapping[str, t.Any]) -> None:
+            schema["example"] = ThermalClusterOutput(
+                id="Gas cluster YZ",
+                group="Gas",
+                name="Gas Cluster YZ",
+                enabled=False,
+                unit_count=100,
+                nominal_capacity=1000.0,
+                gen_ts="use global",
+                co2=7.0,
+            ).model_dump()
+
 
 def create_thermal_output(
     study_version: t.Union[str, int],

diff --git a/antarest/study/business/config_management.py b/antarest/study/business/config_management.py
@@ -42,7 +42,7 @@ def set_playlist(
         file_study = self.storage_service.get_storage(study).get_raw(study)
         command = UpdatePlaylist(
             items=playlist,
-            weights=weights,  # type: ignore
+            weights=weights,
             reverse=reverse,
             active=active,
             command_context=self.storage_service.variant_study_service.command_factory.command_context,

diff --git a/antarest/study/business/correlation_management.py b/antarest/study/business/correlation_management.py
@@ -137,8 +137,19 @@ def validate_correlation_matrix(
 
         return data
 
-
-# TODO SL : restore schema ?
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "columns": ["north", "east", "south", "west"],
+                "data": [
+                    [0.0, 0.0, 0.25, 0.0],
+                    [0.0, 0.0, 0.75, 0.12],
+                    [0.25, 0.75, 0.0, 0.75],
+                    [0.0, 0.12, 0.75, 0.0],
+                ],
+                "index": ["north", "east", "south", "west"],
+            }
+        }
 
 
 def _config_to_array(

diff --git a/antarest/study/business/xpansion_management.py b/antarest/study/business/xpansion_management.py
@@ -240,18 +240,22 @@ class XpansionCandidateDTO(BaseModel):
     name: str
     link: str
     annual_cost_per_mw: float = Field(alias="annual-cost-per-mw", ge=0)
-    unit_size: float = Field(default=None, alias="unit-size", ge=0)
-    max_units: int = Field(default=None, alias="max-units", ge=0)
-    max_investment: float = Field(default=None, alias="max-investment", ge=0)
-    already_installed_capacity: int = Field(default=None, alias="already-installed-capacity", ge=0)
+    unit_size: t.Optional[float] = Field(default=None, alias="unit-size", ge=0)
+    max_units: t.Optional[int] = Field(default=None, alias="max-units", ge=0)
+    max_investment: t.Optional[float] = Field(default=None, alias="max-investment", ge=0)
+    already_installed_capacity: t.Optional[int] = Field(default=None, alias="already-installed-capacity", ge=0)
     # this is obsolete (replaced by direct/indirect)
-    link_profile: str = Field(default=None, alias="link-profile")
+    link_profile: t.Optional[str] = Field(default=None, alias="link-profile")
     # this is obsolete (replaced by direct/indirect)
-    already_installed_link_profile: str = Field(default=None, alias="already-installed-link-profile")
-    direct_link_profile: str = Field(default=None, alias="direct-link-profile")
-    indirect_link_profile: str = Field(default=None, alias="indirect-link-profile")
-    already_installed_direct_link_profile: str = Field(default=None, alias="already-installed-direct-link-profile")
-    already_installed_indirect_link_profile: str = Field(default=None, alias="already-installed-indirect-link-profile")
+    already_installed_link_profile: t.Optional[str] = Field(default=None, alias="already-installed-link-profile")
+    direct_link_profile: t.Optional[str] = Field(default=None, alias="direct-link-profile")
+    indirect_link_profile: t.Optional[str] = Field(default=None, alias="indirect-link-profile")
+    already_installed_direct_link_profile: t.Optional[str] = Field(
+        default=None, alias="already-installed-direct-link-profile"
+    )
+    already_installed_indirect_link_profile: t.Optional[str] = Field(
+        default=None, alias="already-installed-indirect-link-profile"
+    )
 
 
 class LinkNotFound(HTTPException):
@@ -339,7 +343,7 @@ def create_xpansion_configuration(self, study: Study, zipped_config: t.Optional[
                     )
                     raise BadZipBinary("Only zip file are allowed.")
 
-            xpansion_settings = XpansionSettings()  # type: ignore
+            xpansion_settings = XpansionSettings()
             settings_obj = xpansion_settings.model_dump(
                 by_alias=True, exclude_none=True, exclude={"sensitivity_config"}
             )

diff --git a/antarest/study/storage/rawstudy/model/filesystem/config/area.py b/antarest/study/storage/rawstudy/model/filesystem/config/area.py
@@ -104,7 +104,7 @@ class ModalOptimizationSection(IniProperties):
         spread_spilled_energy_cost: float = Field(default=0.0, ge=0, alias="spread-spilled-energy-cost")
 
     filtering: FilteringSection = Field(
-        default_factory=FilteringSection,  # type: ignore
+        default_factory=FilteringSection,
         alias="filtering",
     )
     nodal_optimization: ModalOptimizationSection = Field(

diff --git a/antarest/study/storage/variantstudy/business/utils.py b/antarest/study/storage/variantstudy/business/utils.py
@@ -122,13 +122,13 @@ def transform_command_to_dto(
             commands_dto.append(
                 CommandDTO(
                     action=prev_command.command_name.value,
-                    args=cur_command_args_batch,  # type: ignore
+                    args=cur_command_args_batch,
                 )
             )
             cur_command_args_batch = [command.to_dto().args]
             cur_dto_index += 1
             cur_dto = ref_commands_dto[cur_dto_index]
             cur_dto_arg_count = 1 if isinstance(cur_dto.args, dict) else len(cur_dto.args)
             prev_command = command
-    commands_dto.append(CommandDTO(action=prev_command.command_name.value, args=cur_command_args_batch))  # type: ignore
+    commands_dto.append(CommandDTO(action=prev_command.command_name.value, args=cur_command_args_batch))
     return commands_dto
diff --git a/pyproject.toml b/pyproject.toml
@@ -41,6 +41,7 @@ include = ["antarest*"]
 exclude = "antarest/fastapi_jwt_auth/*"
 strict = true
 files = "antarest"
+plugins = "pydantic.mypy"
 
 [[tool.mypy.overrides]]
 module = ["antarest/fastapi_jwt_auth.*"]