Merge pull request #210 from Limmen/dev6

reward_function_config.py, transistion_operator_config.py
Limmen · Aug 16, 2023 · a55e0c0 · a55e0c0
2 parents 66e35a3 + 081d097
commit a55e0c0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/...n-system/libs/csle-common/src/csle_common/dao/simulation_config/reward_function_config.py b/...n-system/libs/csle-common/src/csle_common/dao/simulation_config/reward_function_config.py
@@ -1,6 +1,7 @@
 from typing import Dict, Any, List
 
 import numpy as np
+from numpy.typing import NDArray
 
 from csle_base.json_serializable import JSONSerializable
 
@@ -10,7 +11,7 @@ class RewardFunctionConfig(JSONSerializable):
     DTO containing the reward tensor of a simulation
     """
 
-    def __init__(self, reward_tensor: List):
+    def __init__(self, reward_tensor: NDArray[Any]):
         """
         Initalizes the DTO
 
@@ -33,8 +34,8 @@ def to_dict(self) -> Dict[str, Any]:
         """
         :return: a dict representation  of the object
         """
-        d = {}
-        if isinstance(self.reward_tensor, np.ndarray):
+        d: Dict[str, Any] = {}
+        if isinstance(self.reward_tensor, type(NDArray[Any])):
             tensor = self.reward_tensor.tolist()
         else:
             tensor = self.reward_tensor

diff --git a/...stem/libs/csle-common/src/csle_common/dao/simulation_config/transition_operator_config.py b/...stem/libs/csle-common/src/csle_common/dao/simulation_config/transition_operator_config.py
@@ -1,5 +1,6 @@
 from typing import List, Dict, Any
 import numpy as np
+from numpy.typing import NDArray
 from csle_base.json_serializable import JSONSerializable
 
 
@@ -8,7 +9,7 @@ class TransitionOperatorConfig(JSONSerializable):
     DTO representing the transition operator definition of a simulation
     """
 
-    def __init__(self, transition_tensor: List):
+    def __init__(self, transition_tensor: NDArray[Any]):
         """
         Initializes the DTO
 
@@ -33,8 +34,8 @@ def to_dict(self) -> Dict[str, Any]:
 
         :return: a dict representation of the object
         """
-        d = {}
-        if isinstance(self.transition_tensor, np.ndarray):
+        d: Dict[str, Any] = {}
+        if isinstance(self.transition_tensor, type(NDArray[Any])):
             tensor = self.transition_tensor.tolist()
         else:
             tensor = self.transition_tensor