allow pd.Timestamp as start, stop

WFP-VAM · Mar 19, 2024 · 367cfeb · 367cfeb
1 parent f80bf11
commit 367cfeb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/hdc/algo/utils.py b/hdc/algo/utils.py
@@ -26,8 +26,8 @@ def to_linspace(x) -> Tuple[np.ndarray, List[int]]:
 
 def get_calibration_indices(
     time: pd.DatetimeIndex,
-    start: str,
-    stop: str,
+    start: Union[str, pd.Timestamp],
+    stop: Union[str, pd.Timestamp],
     groups: Optional[Iterable[Union[int, float, str]]] = None,
     num_groups: Optional[int] = None,
 ) -> Union[Tuple[int, int], np.ndarray]:

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -32,6 +32,8 @@ def test_to_linspace(input_array, expected_output):
 
 def test_get_calibration_indices():
     tix = pd.date_range("2010-01-01", "2010-12-31")
+
+    assert get_calibration_indices(tix, tix[0], tix[-1]) == (0, 365)
     assert get_calibration_indices(tix, "2010-01-01", "2010-12-31") == (0, 365)
     assert get_calibration_indices(tix, "2010-01-15", "2010-12-15") == (14, 349)