typing

Signed-off-by: Sylvain Leclerc <[email protected]>

requirements-dev.txt

@@ -2,4 +2,6 @@ pytest~=7.0.0
 mypy~=1.7.1
 black~=23.7.0
 isort~=5.12.0
-pytest-cov
+pytest-cov
+pandas-stubs
+types-PyYAML

src/cluster_resolve.py

@@ -25,15 +25,15 @@ def resolve_thermal_cluster(
     return ThermalCluster(
         unit_count=parsed_yaml.unit_count,
         nominal_power=parsed_yaml.nominal_power,
-        modulation=modulation["modulation"],
+        modulation=modulation["modulation"].to_list(),
         fo_law=law_dict[parsed_yaml.fo_law],
         fo_volatility=parsed_yaml.fo_volatility,
         po_law=law_dict[parsed_yaml.po_law],
         po_volatility=parsed_yaml.po_volatility,
-        fo_duration=parameters_ts["FOD"],
-        fo_rate=parameters_ts["FOR"],
-        po_duration=parameters_ts["POD"],
-        po_rate=parameters_ts["POR"],
-        npo_min=parameters_ts["POMax"],
-        npo_max=parameters_ts["POMin"],
+        fo_duration=parameters_ts["FOD"].to_list(),
+        fo_rate=parameters_ts["FOR"].to_list(),
+        po_duration=parameters_ts["POD"].to_list(),
+        po_rate=parameters_ts["POR"].to_list(),
+        npo_min=parameters_ts["POMax"].to_list(),
+        npo_max=parameters_ts["POMin"].to_list(),
     )

src/mersenne_twister.py

@@ -38,14 +38,14 @@ class MersenneTwister:
     UPPER_MASK: int = 0x80000000
     LOWER_MASK: int = 0x7FFFFFFF
 
-    MAG: Tuple[int] = (0, MATRIX_A)
+    MAG: Tuple[int, int] = (0, MATRIX_A)
 
     mt: List[int] = [0] * periodN
 
     mti: int = 0
 
     @classmethod
-    def seed(self, seed: int):
+    def seed(self, seed: int) -> None:
         self.mt[0] = seed & 0xFFFFFFFF
         for i in range(1, self.periodN):
             self.mt[i] = 1812433253 * (self.mt[i - 1] ^ (self.mt[i - 1] >> 30)) + i
@@ -54,7 +54,7 @@ def seed(self, seed: int):
         self.mti = self.periodN
 
     @classmethod
-    def next(self):
+    def next(self) -> float:
         if self.mti == self.periodN:
             for j in range(self.periodN - self.periodM):
                 y = (self.mt[j] & self.UPPER_MASK) | (self.mt[j + 1] & self.LOWER_MASK)
@@ -88,5 +88,5 @@ def next(self):
         return y / 4294967295
 
     @classmethod
-    def reset(self):
+    def reset(self) -> None:
         self.seed(5489)

src/ts_generator.py

@@ -18,6 +18,7 @@
 from typing import Any, List, Optional, Tuple
 
 import numpy as np
+from numpy.typing import ArrayLike
 
 from mersenne_twister import MersenneTwister
 

tests/test_parsing.py

@@ -47,14 +47,14 @@ def test(cluster, data_directory):
 
     assert cld.unit_count == cluster.unit_count
     assert cld.nominal_power == cluster.nominal_power
-    assert cld.modulation.to_list() == cluster.modulation
+    assert cld.modulation == cluster.modulation
     assert cld.fo_law == cluster.fo_law
     assert cld.fo_volatility == cluster.fo_volatility
     assert cld.po_law == cluster.po_law
     assert cld.po_volatility == cluster.po_volatility
-    assert cld.fo_duration.to_list() == cluster.fo_duration
-    assert cld.fo_rate.to_list() == cluster.fo_rate
-    assert cld.po_duration.to_list() == cluster.po_duration
-    assert cld.po_rate.to_list() == cluster.po_rate
-    assert cld.npo_min.to_list() == cluster.npo_min
-    assert cld.npo_max.to_list() == cluster.npo_max
+    assert cld.fo_duration == cluster.fo_duration
+    assert cld.fo_rate == cluster.fo_rate
+    assert cld.po_duration == cluster.po_duration
+    assert cld.po_rate == cluster.po_rate
+    assert cld.npo_min == cluster.npo_min
+    assert cld.npo_max == cluster.npo_max