Skip to content

Commit 7c1b09b

Browse files
committed
Use window method in moving window's getitem magic
The `__getitem__` method in moving window was not working properly. Since its functionality is already provided by the window method it one can be called from the getitem magic. Signed-off-by: cwasicki <[email protected]>
1 parent 7a38125 commit 7c1b09b

File tree

2 files changed

+17
-20
lines changed

2 files changed

+17
-20
lines changed

src/frequenz/sdk/timeseries/_moving_window.py

Lines changed: 8 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -367,30 +367,18 @@ def __getitem__(self, key: SupportsIndex | datetime | slice) -> float | ArrayLik
367367
A float if the key is a number or a timestamp.
368368
an numpy array if the key is a slice.
369369
"""
370+
if isinstance(key, slice):
371+
if not (key.step is None or key.step == 1):
372+
raise ValueError("Slicing with a step other than 1 is not supported.")
373+
return self.window(key.start, key.stop)
374+
370375
if self._buffer.count_valid() == 0:
371376
raise IndexError("The buffer is empty.")
372-
if isinstance(key, slice):
373-
if isinstance(key.start, int) or isinstance(key.stop, int):
374-
if key.start is None or key.stop is None:
375-
key = slice(slice(key.start, key.stop).indices(self.count_valid()))
376-
elif isinstance(key.start, datetime) or isinstance(key.stop, datetime):
377-
if key.start is None:
378-
key = slice(self._buffer.time_bound_oldest, key.stop)
379-
if key.stop is None:
380-
key = slice(key.start, self._buffer.time_bound_newest)
381-
382-
_logger.debug("Returning slice for [%s:%s].", key.start, key.stop)
383-
384-
# we are doing runtime typechecks since there is no abstract slice type yet
385-
# see also (https://peps.python.org/pep-0696)
386-
if isinstance(key.start, datetime) and isinstance(key.stop, datetime):
387-
return self._buffer.window(key.start, key.stop)
388-
if isinstance(key.start, int) and isinstance(key.stop, int):
389-
return self._buffer[key]
390-
elif isinstance(key, datetime):
377+
378+
if isinstance(key, datetime):
391379
_logger.debug("Returning value at time %s ", key)
392380
return self._buffer[self._buffer.to_internal_index(key)]
393-
elif isinstance(key, SupportsIndex):
381+
if isinstance(key, SupportsIndex):
394382
_logger.debug("Returning value at index %s ", key)
395383
return self._buffer[key]
396384

tests/timeseries/test_moving_window.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,12 +111,21 @@ async def test_access_window_by_int_slice() -> None:
111111
assert np.array_equal(window[5:14], np.array(data[5:14]))
112112
assert np.array_equal(window.window(5, 14), np.array(data[5:14]))
113113

114+
# Test with step size (other than 1 not supported)
115+
assert np.array_equal(window[5:14:1], np.array(data[5:14]))
116+
assert np.array_equal(window[5:14:None], np.array(data[5:14]))
117+
with pytest.raises(ValueError):
118+
_ = window[5:14:2]
119+
with pytest.raises(ValueError):
120+
_ = window[14:5:-1]
121+
114122
window, sender = init_moving_window(timedelta(seconds=5))
115123

116124
def test_eq(expected: list[float], start: int | None, end: int | None) -> None:
117125
assert np.allclose(
118126
window.window(start, end), np.array(expected), equal_nan=True
119127
)
128+
assert np.allclose(window[start:end], np.array(expected), equal_nan=True)
120129

121130
async with window:
122131
test_eq([], 0, 1)

0 commit comments

Comments
 (0)