Support int indices in ring buffer window method

Using integer-based indices allows for selecting positions within the
window without requiring knowledge of the specific timestamps.

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

Author: cwasicki

src/frequenz/sdk/timeseries/_ringbuffer/buffer.py

@@ -284,8 +284,26 @@ def get_timestamp(self, index: int) -> datetime | None:
         )
         return ref_ts + index * self._sampling_period
 
+    def _to_covered_indices(
+        self, start: int | None, end: int | None = None
+    ) -> tuple[int, int]:
+        """Project the given indices via slice onto the covered range.
+
+        Args:
+            start: Start index.
+            end: End index. Optional, defaults to None.
+
+        Returns:
+            tuple of start and end indices on the range currently covered by the buffer.
+        """
+        return slice(start, end).indices(self.count_covered())[:2]
+
     def window(
-        self, start: datetime, end: datetime, *, force_copy: bool = True
+        self,
+        start: datetime | int | None,
+        end: datetime | int | None,
+        *,
+        force_copy: bool = True,
     ) -> FloatArray:
         """Request a copy or view on the data between start timestamp and end timestamp.
 
@@ -310,17 +328,32 @@ def window(
                 copy of the data.
 
         Raises:
-            IndexError: When requesting a window with invalid timestamps.
+            IndexError: When start and end are not both datetime or index.
 
         Returns:
             The requested window
         """
-        if start > end:
+        if self.count_covered() == 0:
+            return np.array([]) if isinstance(self._buffer, np.ndarray) else []
+
+        # If both are indices or None convert to datetime
+        if not isinstance(start, datetime) and not isinstance(end, datetime):
+            start, end = self._to_covered_indices(start, end)
+            start = self.get_timestamp(start)
+            end = self.get_timestamp(end)
+
+        # Here we should have both as datetime
+        if not isinstance(start, datetime) or not isinstance(end, datetime):
             raise IndexError(
-                f"end parameter {end} has to predate start parameter {start}"
+                f"start ({start}) and end ({end}) must both be either datetime or index."
             )
 
-        if start == end:
+        # Ensure that the window is within the bounds of the buffer
+        assert self.oldest_timestamp is not None and self.newest_timestamp is not None
+        start = max(start, self.oldest_timestamp)
+        end = min(end, self.newest_timestamp + self._sampling_period)
+
+        if start >= end:
             return np.array([]) if isinstance(self._buffer, np.ndarray) else []
 
         start_pos = self.to_internal_index(start)

tests/timeseries/test_moving_window.py

@@ -121,13 +121,10 @@ async def test_access_window_by_ts_slice() -> None:
         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))
+        # Window also supports slicing with indices outside allowed range
+        assert np.array_equal(window.window(dt(3), dt(1)), np.array([]))
+        assert np.array_equal(window.window(dt(3), dt(6)), np.array([3, 4]))
+        assert np.array_equal(window.window(dt(-1), dt(5)), np.array([0, 1, 2, 3, 4]))
 
 
 async def test_access_empty_window() -> None:

tests/timeseries/test_ringbuffer.py

@@ -137,12 +137,11 @@ def test_timestamp_ringbuffer_gaps(
         Sample(datetime.fromtimestamp(500 + size, tz=timezone.utc), Quantity(9999))
     )
 
-    # Expect exception for the same window
-    with pytest.raises(IndexError):
-        buffer.window(
-            datetime.fromtimestamp(200, tz=timezone.utc),
-            datetime.fromtimestamp(202, tz=timezone.utc),
-        )
+    # Allow still to request old (empty) window
+    buffer.window(
+        datetime.fromtimestamp(200, tz=timezone.utc),
+        datetime.fromtimestamp(202, tz=timezone.utc),
+    )
 
     # Receive new window without exception
     buffer.window(
@@ -524,8 +523,8 @@ def get_orb(data: FloatArray) -> OrderedRingBuffer[FloatArray]:
     return buffer
 
 
-def test_window() -> None:
-    """Test the window function."""
+def test_window_datetime() -> None:
+    """Test the window function with datetime."""
     buffer = get_orb(np.array([0, None, 2, 3, 4]))
     win = buffer.window(dt(0), dt(3), force_copy=False)
     assert [0, np.nan, 2] == list(win)
@@ -543,6 +542,39 @@ def test_window() -> None:
     assert [] == buffer.window(dt(1), dt(1))
 
 
+def test_window_index() -> None:
+    """Test the window function with index."""
+    buffer = get_orb([0.0, 1.0, 2.0, 3.0, 4.0])
+    assert [0, 1, 2] == buffer.window(0, 3)
+    assert [0, 1, 2, 3, 4] == buffer.window(0, 5)
+    assert [0, 1, 2, 3, 4] == buffer.window(0, 99)
+    assert [2, 3] == buffer.window(-3, -1)
+    assert [2, 3, 4] == buffer.window(-3, 5)
+    assert [0, 1, 2, 3] == buffer.window(-5, -1)
+    assert [0, 1, 2, 3, 4] == buffer.window(-99, None)
+    assert [0, 1, 2, 3, 4] == buffer.window(None, 99)
+    # start >= end
+    assert [] == buffer.window(0, 0)
+    assert [] == buffer.window(-5, 0)
+    assert [] == buffer.window(1, 0)
+    assert [] == buffer.window(-1, -2)
+    assert [] == buffer.window(-3, 0)
+
+
+def test_window_fail() -> None:
+    """Test the window function with invalid indices."""
+    buffer = get_orb([0.0, 1.0, 2.0, 3.0, 4.0])
+    # Go crazy with the indices
+    with pytest.raises(IndexError):
+        buffer.window(dt(1), 3)
+    with pytest.raises(IndexError):
+        buffer.window(1, dt(3))
+    with pytest.raises(IndexError):
+        buffer.window(None, dt(2))
+    with pytest.raises(IndexError):
+        buffer.window(dt(2), None)
+
+
 def test_wrapped_buffer_window() -> None:
     """Test the wrapped buffer window function."""
     wbw = OrderedRingBuffer._wrapped_buffer_window  # pylint: disable=protected-access