Add at method for single element access of moving window

This method implements corrected access of single elements in the moving
window. Similar to its `window` counterpart for slices, in future it
could provide additional features such as replacing missing values.

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

Author: cwasicki

src/frequenz/sdk/timeseries/_moving_window.py

@@ -241,6 +241,48 @@ def capacity(self) -> int:
         """
         return self._buffer.maxlen
 
+    # pylint before 3.0 only accepts names with 3 or more chars
+    def at(self, key: int | datetime) -> float:  # pylint: disable=invalid-name
+        """
+        Return the sample at the given index or timestamp.
+
+        In contrast to the [`window`][frequenz.sdk.timeseries.MovingWindow.window] method,
+        which expects a slice as argument, this method expects a single index as argument
+        and returns a single value.
+
+        Args:
+            key: The index or timestamp of the sample to return.
+
+        Returns:
+            The sample at the given index or timestamp.
+
+        Raises:
+            IndexError: If the buffer is empty or the index is out of bounds.
+        """
+        if self._buffer.count_valid() == 0:
+            raise IndexError("The buffer is empty.")
+
+        if isinstance(key, datetime):
+            assert self._buffer.oldest_timestamp is not None
+            assert self._buffer.newest_timestamp is not None
+            if (
+                key < self._buffer.oldest_timestamp
+                or key > self._buffer.newest_timestamp
+            ):
+                raise IndexError(
+                    f"Timestamp {key} is out of range [{self._buffer.oldest_timestamp}, "
+                    f"{self._buffer.newest_timestamp}]"
+                )
+            return self._buffer[self._buffer.to_internal_index(key)]
+
+        if isinstance(key, int):
+            _logger.debug("Returning value at index %s ", key)
+            timestamp = self._buffer.get_timestamp(key)
+            assert timestamp is not None
+            return self._buffer[self._buffer.to_internal_index(timestamp)]
+
+        raise TypeError("Key has to be either a timestamp or an integer.")
+
     def window(
         self,
         start: datetime | int | None,
@@ -251,6 +293,10 @@ def window(
         """
         Return an array containing the samples in the given time interval.
 
+        In contrast to the [`at`][frequenz.sdk.timeseries.MovingWindow.at] method,
+        which expects a single index as argument, this method expects a slice as argument
+        and returns an array.
+
         Args:
             start: The start of the time interval. If `None`, the start of the
                 window is used.

tests/timeseries/test_moving_window.py

@@ -99,12 +99,18 @@ async def test_access_window_by_timestamp() -> None:
     async with window:
         await push_logical_meter_data(sender, [0, 1, 2])
         assert np.array_equal(window[dt(1)], 1.0)
+        assert np.array_equal(window.at(dt(1)), 1.0)
         assert np.array_equal(window[dt(2)], 2.0)
+        assert np.array_equal(window.at(dt(2)), 2.0)
         assert np.array_equal(window[dt(3)], 1.0)  # bug: should raise
         with pytest.raises(IndexError):
             _ = window[dt(0)]
+        with pytest.raises(IndexError):
+            _ = window.at(dt(0))
         with pytest.raises(IndexError):
             _ = window[dt(4)]
+        with pytest.raises(IndexError):
+            _ = window.at(dt(3))
 
 
 async def test_access_window_by_int_slice() -> None: