change type hint

Signed-off-by: Nitish Bharambe <nitish.bharambe@alliander.com>
PowerGridModel · Jul 22, 2024 · ba7ce89 · ba7ce89
1 parent 3ec562a
commit ba7ce89
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/src/power_grid_model/_utils.py b/src/power_grid_model/_utils.py
@@ -294,7 +294,7 @@ def copy_output_to_columnar_dataset(
     output_component_types: ComponentAttributeMapping,
     output_type: OutputType,
     available_components: list[ComponentType],
-) -> Dataset:
+) -> dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
     """Temporary function to copy row based dataset to a column based dataset as per output_component_types.
     The purpose of this function is to mimic columnar data without any memory overhead benefits.
 

diff --git a/src/power_grid_model/core/power_grid_model.py b/src/power_grid_model/core/power_grid_model.py
@@ -311,7 +311,7 @@ def _calculate_power_flow(
         decode_error: bool = True,
         tap_changing_strategy: Union[TapChangingStrategy, str] = TapChangingStrategy.disabled,
         experimental_features: Union[_ExperimentalFeatures, str] = _ExperimentalFeatures.disabled,
-    ):
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         calculation_type = CalculationType.power_flow
         options = self._options(
             calculation_type=calculation_type,
@@ -346,7 +346,7 @@ def _calculate_state_estimation(
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         experimental_features: Union[_ExperimentalFeatures, str] = _ExperimentalFeatures.disabled,
-    ) -> Dict[ComponentType, np.ndarray]:
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         calculation_type = CalculationType.state_estimation
         options = self._options(
             calculation_type=calculation_type,
@@ -378,7 +378,7 @@ def _calculate_short_circuit(
         decode_error: bool = True,
         short_circuit_voltage_scaling: Union[ShortCircuitVoltageScaling, str] = ShortCircuitVoltageScaling.maximum,
         experimental_features: Union[_ExperimentalFeatures, str] = _ExperimentalFeatures.disabled,
-    ) -> Dict[ComponentType, np.ndarray]:
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         calculation_type = CalculationType.short_circuit
         symmetric = False
 
@@ -418,7 +418,7 @@ def calculate_power_flow(
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         tap_changing_strategy: Union[TapChangingStrategy, str] = TapChangingStrategy.disabled,
-    ) -> Dict[ComponentType, np.ndarray]:
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         """
         Calculate power flow once with the current model attributes.
         Or calculate in batch with the given update dataset in batch.
@@ -511,7 +511,7 @@ def calculate_state_estimation(
         output_component_types: Optional[Union[Set[ComponentType], List[ComponentType]]] = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
-    ) -> Dict[ComponentType, np.ndarray]:
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         """
         Calculate state estimation once with the current model attributes.
         Or calculate in batch with the given update dataset in batch.
@@ -598,7 +598,7 @@ def calculate_short_circuit(
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         short_circuit_voltage_scaling: Union[ShortCircuitVoltageScaling, str] = ShortCircuitVoltageScaling.maximum,
-    ) -> Dict[ComponentType, np.ndarray]:
+    ) -> Dict[ComponentType, np.ndarray] | dict[ComponentType, dict[str, np.ndarray]]:
         """
         Calculate a short circuit once with the current model attributes.
         Or calculate in batch with the given update dataset in batch