diff --git a/quixstreams/dataframe/windows/base.py b/quixstreams/dataframe/windows/base.py index 1c4b69aaa..696f64122 100644 --- a/quixstreams/dataframe/windows/base.py +++ b/quixstreams/dataframe/windows/base.py @@ -1,4 +1,5 @@ -from typing import Any, Optional, Callable, Tuple, List +from collections import deque +from typing import Any, Deque, Optional, Callable, Tuple from typing_extensions import TypedDict @@ -17,7 +18,7 @@ class WindowResult(TypedDict): def get_window_ranges( timestamp_ms: int, duration_ms: int, step_ms: Optional[int] = None -) -> List[Tuple[int, int]]: +) -> Deque[Tuple[int, int]]: """ Get a list of window ranges for the given timestamp. :param timestamp_ms: timestamp in milliseconds @@ -28,14 +29,14 @@ def get_window_ranges( if not step_ms: step_ms = duration_ms - window_ranges = [] + window_ranges = deque() current_window_start = timestamp_ms - (timestamp_ms % step_ms) while ( current_window_start > timestamp_ms - duration_ms and current_window_start >= 0 ): window_end = current_window_start + duration_ms - window_ranges.insert(0, (current_window_start, window_end)) + window_ranges.appendleft((current_window_start, window_end)) current_window_start -= step_ms return window_ranges diff --git a/tests/test_quixstreams/test_dataframe/test_windows/test_base.py b/tests/test_quixstreams/test_dataframe/test_windows/test_base.py index 75cb57c9a..36c7f11e8 100644 --- a/tests/test_quixstreams/test_dataframe/test_windows/test_base.py +++ b/tests/test_quixstreams/test_dataframe/test_windows/test_base.py @@ -16,12 +16,10 @@ class TestGetWindowRanges: ], ) def test_get_window_ranges_with_step(self, timestamp, duration, step, expected): - assert ( - get_window_ranges( - timestamp_ms=timestamp, duration_ms=duration, step_ms=step - ) - == expected + result = get_window_ranges( + timestamp_ms=timestamp, duration_ms=duration, step_ms=step ) + assert list(result) == expected @pytest.mark.parametrize( "timestamp, duration, expected", @@ -31,4 +29,5 @@ def test_get_window_ranges_with_step(self, timestamp, duration, step, expected): ], ) def test_get_window_ranges_no_step(self, timestamp, duration, expected): - assert get_window_ranges(timestamp, duration) == expected + result = get_window_ranges(timestamp, duration) + assert list(result) == expected