|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | | -from typing import Iterable, TypeVar |
| 3 | +from typing import Iterable, Literal, Sequence, TypeVar |
4 | 4 |
|
5 | 5 | T = TypeVar("T") |
6 | 6 |
|
@@ -43,3 +43,44 @@ def loop_first_last(values: Iterable[T]) -> Iterable[tuple[bool, bool, T]]: |
43 | 43 | first = False |
44 | 44 | previous_value = value |
45 | 45 | yield first, True, previous_value |
| 46 | + |
| 47 | + |
| 48 | +def loop_from_index( |
| 49 | + values: Sequence[T], |
| 50 | + index: int, |
| 51 | + direction: Literal[-1, +1] = +1, |
| 52 | + wrap: bool = True, |
| 53 | +) -> Iterable[tuple[int, T]]: |
| 54 | + """Iterate over values in a sequence from a given starting index, potentially wrapping the index |
| 55 | + if it would go out of bounds. |
| 56 | +
|
| 57 | + Note that the first value to be yielded is a step from `index`, and `index` will be yielded *last*. |
| 58 | +
|
| 59 | +
|
| 60 | + Args: |
| 61 | + values: A sequence of values. |
| 62 | + index: Starting index. |
| 63 | + direction: Direction to move index (+1 for forward, -1 for backward). |
| 64 | + bool: Should the index wrap when out of bounds? |
| 65 | +
|
| 66 | + Yields: |
| 67 | + A tuple of index and value from the sequence. |
| 68 | + """ |
| 69 | + # Sanity check for devs who miss the typing errors |
| 70 | + assert direction in (-1, +1), "direction must be -1 or +1" |
| 71 | + count = len(values) |
| 72 | + if wrap: |
| 73 | + for _ in range(count): |
| 74 | + index = (index + direction) % count |
| 75 | + yield (index, values[index]) |
| 76 | + else: |
| 77 | + if direction == +1: |
| 78 | + for _ in range(count): |
| 79 | + if (index := index + 1) >= count: |
| 80 | + break |
| 81 | + yield (index, values[index]) |
| 82 | + else: |
| 83 | + for _ in range(count): |
| 84 | + if (index := index - 1) < 0: |
| 85 | + break |
| 86 | + yield (index, values[index]) |
0 commit comments