diff --git a/boltons/iterutils.py b/boltons/iterutils.py index 3707566d..a8d52c3c 100644 --- a/boltons/iterutils.py +++ b/boltons/iterutils.py @@ -41,12 +41,16 @@ following are based on examples in itertools docs. """ +from __future__ import annotations +import dataclasses import os import math import time import codecs import random import itertools +import operator +from typing import Iterable try: from collections.abc import Mapping, Sequence, Set, ItemsView, Iterable @@ -1510,3 +1514,145 @@ def __lt__(self, other): $ python -m timeit -s "x = [1]" "try: x.split('.') \nexcept AttributeError: pass" 1000000 loops, best of 3: 0.544 usec per loop """ + + +def _iter_or_num(thing: Union[Iterable[float], float]) -> Iterator[float]: + if isinstance(thing, (int, float)): + return itertools.repeat(thing) + return iter(thing) + + +@dataclasses.dataclass +class NumIterator: + + """An iterator of numbers. + + Supports math operations and slicing. + + >>> list(NumIterator(range(3))) + [0, 1, 2] + >>> list(NumIterator(range(3)) + 1) + [1, 2, 3] + >>> list(NumIterator(range(3)) * 2) + [0, 2, 4] + >>> list(NumIterator(x + 1 for x in range(10))[:3]) + [1, 2, 3] + + There are also a few helper class methods to generate + common iterators. + They all generate *infinite* iterators. + Convert them to finite iterators using slicing. + + Counting: + + >>> list(NumIterator.count()[:5]) + [0, 1, 2, 3, 4] + + Constant: + + >>> list(NumIterator.constant(7)[:2]) + [7, 7] + + Fibonacci sequence: + + >>> list(NumIterator.fib()[:8]) + [0, 1, 1, 2, 3, 5, 8, 13] + + Since the iterators support math operations, + you can combine them to generate more sophisticated + sequences. + + For example, + + >>> fib_pow_of_2 = 2 ** NumIterator.count() + NumIterator.fib() + >>> list(fib_pow_of_2[:5]) + [1, 3, 5, 10, 19] + """ + + _original: Iterable[float] + + @classmethod + def count(cls): + return cls(itertools.count()) + + @classmethod + def fib(cls): + def inner_fib(): + a, b = 0, 1 + yield a + while True: + yield b + a, b = b, a+b + return cls(inner_fib()) + + @classmethod + def constant(cls, num): + return cls(itertools.repeat(num)) + + def __iter__(self) -> Iterator[float]: + return iter(self._original) + + def apply_operator(self, op: Callable[[float, float], float], other: Union[Iterable[float], float]) -> NumIterator: + return NumIterator(map(op, self._original, _iter_or_num(other))) + + def apply_r_operator(self, op: Callable[[float, float], float], other: Union[Iterable[float], float]) -> NumIterator: + return NumIterator(map(op, _iter_or_num(other), self._original)) + + def __getitem__(self, a_slice: slice)-> NumIterator: + if not isinstance(a_slice, slice): + raise TypeError( + "no random access, can only slice", + a_slice, + ) + new_original = itertools.islice( + self._original, + a_slice.start, + a_slice.stop, + a_slice.step, + ) + return NumIterator(new_original) + + def __add__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.add, other) + + def __radd__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.add, other) + + def __mul__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.mul, other) + + def __rmul__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.mul, other) + + def __floordiv__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.floordiv, other) + + def __rfloordiv__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.floordiv, other) + + def __truediv__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.truediv, other) + + def __rtruediv__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.truediv, other) + + def __sub__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.sub, other) + + def __rsub__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.sub, other) + + def __pow__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_operator(operator.pow, other) + + def __rpow__(self, other: Union[Iterable[float], float]) -> NumIterator: + return self.apply_r_operator(operator.pow, other) + + def __neg__(self, *args) -> NumIterator: + return NumIterator(map(operator.neg, self)) + + def __pos__(self, *args) -> NumIterator: + return NumIterator(map(operator.pos, self)) + + def __round__(self, *args) -> NumIterator: + return NumIterator((round(item, *args) for item in self)) diff --git a/docs/iterutils.rst b/docs/iterutils.rst index 1cb5c759..0e22e1ab 100644 --- a/docs/iterutils.rst +++ b/docs/iterutils.rst @@ -115,3 +115,12 @@ In the same vein as the feature-checking builtin, :func:`callable`. .. autofunction:: is_iterable .. autofunction:: is_scalar .. autofunction:: is_collection + +Numeric Iterators +----------------- + +A class to wrap iterators producing numbers +to allow it to support arithmetic operations +and slicing. + +.. autoclass:: NumIterator \ No newline at end of file diff --git a/tests/test_iterutils.py b/tests/test_iterutils.py index 0896f38a..de2f8b11 100644 --- a/tests/test_iterutils.py +++ b/tests/test_iterutils.py @@ -9,7 +9,9 @@ research, default_enter, default_exit, - get_path) + get_path, + NumIterator, + ) from boltons.namedutils import namedtuple CUR_PATH = os.path.abspath(__file__) @@ -535,3 +537,64 @@ def test_strip(): assert strip([0,0,0,1,0,2,0,3,0,0,0],0) == [1,0,2,0,3] assert strip([]) == [] +class TestNumIterator: + def test_constant(self): + assert list(NumIterator.constant(5)[:3]) == [5, 5, 5] + + def test_fib(self): + assert list(NumIterator.fib()[:4]) == [0, 1, 1, 2] + + def test_count(self): + assert list(NumIterator.count()[:4]) == [0, 1, 2, 3] + + def test_slice(self): + assert list(NumIterator.count()[1:4:2]) == [1, 3] + + def test_no_access(self): + with pytest.raises(TypeError): + NumIterator.count()[5] + + def test_add(self): + assert list(NumIterator.count()[:3] + 1) == [1, 2, 3] + + def test_radd(self): + assert list(1 + NumIterator.count()[:3]) == [1, 2, 3] + + def test_sub(self): + assert list(NumIterator.count()[:3] - 1) == [-1, 0, 1] + + def test_rsub(self): + assert list(2 - NumIterator.count()[:3]) == [2, 1, 0] + + def test_mul(self): + assert list(NumIterator.count()[:3] * 2) == [0, 2, 4] + + def test_rmul(self): + assert list(2 * NumIterator.count()[:3]) == [0, 2, 4] + + def test_pow(self): + assert list(NumIterator.count()[:3] ** 2) == [0, 1, 4] + + def test_rpow(self): + assert list(2 ** NumIterator.count()[:3]) == [1, 2, 4] + + def test_div(self): + assert list(NumIterator.count()[:2] / 2) == [0.0, 0.5] + + def test_floordiv(self): + assert list(NumIterator.count()[:3] // 2) == [0, 0, 1] + + def test_rdiv(self): + assert list(1 / NumIterator.count()[1:3]) == [1.0, 0.5] + + def test_rfloordiv(self): + assert list(2 // NumIterator.count()[1:5]) == [2, 1, 0, 0] + + def test_neg(self): + assert list(-NumIterator.count()[:3]) == [0, -1, -2] + + def test_pos(self): + assert list(+NumIterator.count()[:3]) == [0, +1, +2] + + def test_pos(self): + assert list(round(NumIterator.count()[:4] / 3)) == [0, 0, 1, 1]