Support int indices for window method in moving window

Signed-off-by: cwasicki <[email protected]>

Author: cwasicki

src/frequenz/sdk/timeseries/_moving_window.py

@@ -243,30 +243,26 @@ def capacity(self) -> int:
 
     def window(
         self,
-        start: datetime,
-        end: datetime,
+        start: datetime | int | None,
+        end: datetime | int | None,
         *,
         force_copy: bool = True,
     ) -> ArrayLike:
         """
         Return an array containing the samples in the given time interval.
 
         Args:
-            start: The start of the time interval. Only datetime objects are supported.
-            end: The end of the time interval. Only datetime objects are supported.
+            start: The start of the time interval. If `None`, the start of the
+                window is used.
+            end: The end of the time interval. If `None`, the end of the window
+                is used.
             force_copy: If `True`, the returned array is a copy of the underlying
                 data. Otherwise, if possible, a view of the underlying data is
                 returned.
 
         Returns:
             An array containing the samples in the given time interval.
-
-        Raises:
-            IndexError: if `start` or `end` are not datetime objects.
         """
-        if not isinstance(start, datetime) or not isinstance(end, datetime):
-            raise IndexError("Only datetime objects are supported as start and end.")
-
         return self._buffer.window(start, end, force_copy=force_copy)
 
     async def _run_impl(self) -> None:

tests/timeseries/test_moving_window.py

@@ -104,11 +104,52 @@ async def test_access_window_by_int_slice() -> None:
     async with window:
         await push_logical_meter_data(sender, range(0, 5))
         assert np.array_equal(window[3:5], np.array([3.0, 4.0]))
-        with pytest.raises(IndexError):
-            window.window(3, 5)  # type: ignore
+        assert np.array_equal(window.window(3, 5), np.array([3.0, 4.0]))
+
         data = [1, 2, 2.5, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1]
         await push_logical_meter_data(sender, data)
         assert np.array_equal(window[5:14], np.array(data[5:14]))
+        assert np.array_equal(window.window(5, 14), np.array(data[5:14]))
+
+    window, sender = init_moving_window(timedelta(seconds=5))
+
+    def test_eq(expected: list[float], start: int | None, end: int | None) -> None:
+        assert np.allclose(
+            window.window(start, end), np.array(expected), equal_nan=True
+        )
+
+    async with window:
+        test_eq([], 0, 1)
+
+        # Incomplete window
+        await push_logical_meter_data(sender, [0.0, 1.0])
+        test_eq([0.0, 1.0], 0, 2)
+        test_eq([0.0, 1.0], 0, 9)
+        test_eq([0.0, 1.0], 0, None)
+        test_eq([0.0, 1.0], -9, None)
+        test_eq([0.0, 1.0], None, None)
+        test_eq([0.0], -2, -1)
+        test_eq([1.0], -1, None)
+
+        # Incomplete window with gap
+        await push_logical_meter_data(
+            sender, [3.0], start_ts=UNIX_EPOCH + timedelta(seconds=3)
+        )
+        test_eq([0.0, 1.0], 0, 2)
+        # gap fill not supported yet:
+        # test_eq([0.0, 1.0, np.nan, 3.0], 0, None)
+        # test_eq([0.0, 1.0, np.nan, 3.0], -9, None)
+        # test_eq([np.nan, 3.0], -2, None)
+
+        # Complete window
+        await push_logical_meter_data(sender, [0.0, 1.0, 2.0, 3.0, 4.0])
+        test_eq([0.0, 1.0], 0, 2)
+        test_eq([3.0, 4.0], -2, None)
+
+        # Complete window with nan
+        await push_logical_meter_data(sender, [0.0, 1.0, np.nan])
+        test_eq([0.0, 1.0, np.nan], 0, 3)
+        test_eq([np.nan, 3.0, 4.0], -3, None)
 
 
 async def test_access_window_by_ts_slice() -> None: