NumIterator

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))

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

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]