Add window function to moving window

Expose window functionality of underlying buffer. Only datetime indices
are supported so far.

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

Author: cwasicki

src/frequenz/sdk/timeseries/_moving_window.py

@@ -241,6 +241,34 @@ def capacity(self) -> int:
         """
         return self._buffer.maxlen
 
+    def window(
+        self,
+        start: datetime,
+        end: datetime,
+        *,
+        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.
+            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:
         """Awaits samples from the receiver and updates the underlying ring buffer.
 

tests/timeseries/test_moving_window.py

@@ -66,6 +66,18 @@ def init_moving_window(
     return window, lm_tx
 
 
+def dt(i: int) -> datetime:  # pylint: disable=invalid-name
+    """Create datetime objects from indices.
+
+    Args:
+        i: Index to create datetime from.
+
+    Returns:
+        Datetime object.
+    """
+    return datetime.fromtimestamp(i, tz=timezone.utc)
+
+
 async def test_access_window_by_index() -> None:
     """Test indexing a window by integer index."""
     window, sender = init_moving_window(timedelta(seconds=1))
@@ -92,7 +104,8 @@ 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
         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]))
@@ -106,6 +119,15 @@ async def test_access_window_by_ts_slice() -> None:
         time_start = UNIX_EPOCH + timedelta(seconds=3)
         time_end = time_start + timedelta(seconds=2)
         assert np.array_equal(window[time_start:time_end], np.array([3.0, 4.0]))  # type: ignore
+        assert np.array_equal(window.window(dt(3), dt(5)), np.array([3.0, 4.0]))
+        assert np.array_equal(window.window(dt(3), dt(3)), np.array([]))
+        # Window only supports slicing with ascending indices within allowed range
+        with pytest.raises(IndexError):
+            window.window(dt(3), dt(1))
+        with pytest.raises(IndexError):
+            window.window(dt(3), dt(6))
+        with pytest.raises(IndexError):
+            window.window(dt(-1), dt(5))
 
 
 async def test_access_empty_window() -> None: