Skip to content

Commit

Permalink
feat: ensure Iter is an iterator.
Browse files Browse the repository at this point in the history
Fixes #71
  • Loading branch information
MartinBernstorff committed Jan 4, 2024
1 parent 87298aa commit 124d807
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 11 deletions.
18 changes: 9 additions & 9 deletions iterpy/_iter.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,22 +17,22 @@

class Iter(Generic[T]):
def __init__(self, iterable: Iterable[T]) -> None:
self._iterator: Iterator[T] = iter(iterable)
self.__consumable_iterator: Iterator[T] = iter(iterable)

@property
def _non_consumable(self) -> Iterator[T]:
return deepcopy(self._iterator)
def _iterator(self) -> Iterator[T]:
return deepcopy(self.__consumable_iterator)

def __getitem__(self, index: int | slice) -> T | "Iter[T]":
if isinstance(index, int) and index >= 0:
try:
return list(self._non_consumable)[index]
return list(self._iterator)[index]
except StopIteration:
raise IndexError("Index out of range") from None
elif isinstance(index, slice):
return Iter(
islice(
self._non_consumable,
self._iterator,
index.start,
index.stop,
index.step,
Expand All @@ -44,6 +44,9 @@ def __getitem__(self, index: int | slice) -> T | "Iter[T]":
)

### Reductions
def reduce(self, func: Callable[[T, T], T]) -> T:
return reduce(func, self._iterator)

def count(self) -> int:
return sum(1 for _ in self._iterator)

Expand All @@ -54,7 +57,7 @@ def to_list(self) -> list[T]:
def to_tuple(self) -> tuple[T, ...]:
return tuple(self._iterator) # pragma: no cover

def to_iterator(self) -> Iterator[T]:
def to_consumable(self) -> Iterator[T]:
return iter(self._iterator) # pragma: no cover

def to_set(self) -> set[T]:
Expand All @@ -81,9 +84,6 @@ def pmap(
def filter(self, func: Callable[[T], bool]) -> "Iter[T]": # noqa: A003
return Iter(filter(func, self._iterator))

def reduce(self, func: Callable[[T, T], T]) -> T:
return reduce(func, self._iterator)

def groupby(
self, func: Callable[[T], str]
) -> "Iter[tuple[str, list[T]]]":
Expand Down
2 changes: 1 addition & 1 deletion iterpy/_iter.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ class Iter(Generic[T]):
def count(self) -> int: ...
def to_list(self) -> list[T]: ...
def to_tuple(self) -> tuple[T, ...]: ...
def to_iterator(self) -> Iterator[T]: ...
def to_consumable(self) -> Iterator[T]: ...
def to_set(self) -> set[T]: ...
def map(self, func: Callable[[T], S]) -> Iter[S]: # noqa: A003
...
Expand Down
2 changes: 1 addition & 1 deletion iterpy/benchmark/query_1/iterators_q1.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ def main_iterator(data: Sequence[Item]) -> Sequence[CategorySummary]:

summaries = [
summarise_category(group, rows=values)
for group, values in mapping.to_iterator()
for group, values in mapping.to_consumable()
]

return summaries
Expand Down
12 changes: 12 additions & 0 deletions iterpy/test_iter.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,18 @@ def test_getitem():
assert test_iterator[0:2].to_list() == [1, 2]


def test_iteration():
test_iterator = Iter([1, 2, 3])
for i in test_iterator.to_consumable():
assert i in [1, 2, 3]


def test_statefulness():
test_iterator = Iter([1, 2, 3])
assert test_iterator.to_list() == [1, 2, 3]
assert test_iterator.to_list() == [1, 2, 3]


def test_flatten():
test_input: list[list[int]] = [[1, 2], [3, 4]]
iterator = Iter(test_input)
Expand Down

0 comments on commit 124d807

Please sign in to comment.