From 1e96a4f9beba18762aa0a428602ffc3b00d2229a Mon Sep 17 00:00:00 2001
From: Frode Helgetun Krogh <70878501+frodehk@users.noreply.github.com>
Date: Wed, 15 May 2024 15:47:21 +0200
Subject: [PATCH] fix: use custom type pydantic class to wrap ndarray and avoid
 validation problem (#487) (#488)

* fix: use custom type pydantic class to wrap ndarray and avoid validation problems

ECALC-965

(cherry picked from commit ab2877ecf68802b436347ed4295bedc9fc9cb705)
---
 .../consumers/legacy_consumer/system/results.py     | 13 +++++++------
 src/libecalc/expression/expression_evaluator.py     |  3 ++-
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/libecalc/core/consumers/legacy_consumer/system/results.py b/src/libecalc/core/consumers/legacy_consumer/system/results.py
index e5e60f9c19..37ded4ec14 100644
--- a/src/libecalc/core/consumers/legacy_consumer/system/results.py
+++ b/src/libecalc/core/consumers/legacy_consumer/system/results.py
@@ -19,6 +19,7 @@
 )
 from libecalc.core.models.results import CompressorTrainResult, PumpModelResult
 from libecalc.core.result.results import ConsumerModelResult
+from libecalc.core.utils.array_type import PydanticNDArray
 
 
 class ConsumerSystemComponentResult(BaseModel):
@@ -30,7 +31,7 @@ def energy_usage(self) -> List[Optional[float]]:
         return self.consumer_model_result.energy_usage
 
     @property
-    def power(self) -> NDArray[np.float64]:
+    def power(self) -> PydanticNDArray:
         if self.consumer_model_result.power is not None:
             return np.asarray(self.consumer_model_result.power)
         else:
@@ -58,7 +59,7 @@ class ConsumerSystemOperationalSettingResult(BaseModel):
     model_config = ConfigDict(frozen=True)
 
     @property
-    def total_energy_usage(self) -> NDArray[np.float64]:
+    def total_energy_usage(self) -> PydanticNDArray:
         total_energy_usage = np.sum(
             [np.asarray(result.energy_usage) for result in self.consumer_results],
             axis=0,
@@ -66,7 +67,7 @@ def total_energy_usage(self) -> NDArray[np.float64]:
         return np.array(total_energy_usage)
 
     @property
-    def total_power(self) -> NDArray[np.float64]:
+    def total_power(self) -> PydanticNDArray:
         total_power = np.sum(
             [np.asarray(result.power) for result in self.consumer_results],
             axis=0,
@@ -74,7 +75,7 @@ def total_power(self) -> NDArray[np.float64]:
         return np.array(total_power)
 
     @property
-    def indices_outside_capacity(self) -> NDArray[np.float64]:
+    def indices_outside_capacity(self) -> PydanticNDArray:
         invalid_indices = np.full_like(self.total_energy_usage, fill_value=0)
 
         for result in self.consumer_results:
@@ -110,12 +111,12 @@ class ConsumerSystemConsumerFunctionResult(ConsumerFunctionResultBase):
 
     typ: Literal[ConsumerFunctionType.SYSTEM] = ConsumerFunctionType.SYSTEM  # type: ignore[valid-type]
 
-    operational_setting_used: NDArray[np.float64]  # integers in the range of number of operational settings
+    operational_setting_used: PydanticNDArray  # integers in the range of number of operational settings
     operational_settings: List[List[ConsumerSystemOperationalSetting]]
     operational_settings_results: List[List[ConsumerSystemOperationalSettingResult]]
     consumer_results: List[List[ConsumerSystemComponentResult]]
     cross_over_used: Optional[
-        NDArray[np.float64]
+        PydanticNDArray
     ] = None  # 0 or 1 whether cross over is used for this result (1=True, 0=False)
 
     def extend(self, other) -> ConsumerSystemConsumerFunctionResult:
diff --git a/src/libecalc/expression/expression_evaluator.py b/src/libecalc/expression/expression_evaluator.py
index cb8a2abc54..892d29db25 100644
--- a/src/libecalc/expression/expression_evaluator.py
+++ b/src/libecalc/expression/expression_evaluator.py
@@ -12,6 +12,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 from libecalc.common.logger import logger
+from libecalc.core.utils.array_type import PydanticNDArray
 
 """
 Module for expression parsing used in Energy/CO2/emissions calculator
@@ -405,7 +406,7 @@ class Operators(Enum):
 
 class Token(BaseModel):
     tag: TokenTag
-    value: Union[float, int, bool, NDArray[np.float64], str] = Field(union_mode="left_to_right")
+    value: Union[float, int, bool, PydanticNDArray, str] = Field(union_mode="left_to_right")
 
     def __str__(self):
         return str(self.value)