New utility to cache and re-access iterators

Author: carver

eth/_utils/generator.py

@@ -0,0 +1,25 @@
+import itertools
+from typing import (
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    TypeVar,
+)
+
+
+TItem = TypeVar('TItem')
+
+
+class CachedIterable(Generic[TItem], Iterable[TItem]):
+    def __init__(self, iterator: Iterable[TItem]) -> None:
+        self._cached_results: List[TItem] = []
+        self._iterator = iterator
+
+    def __iter__(self) -> Iterator[TItem]:
+        return itertools.chain(self._cached_results, self._cache_and_yield())
+
+    def _cache_and_yield(self) -> Iterable[TItem]:
+        for item in self._iterator:
+            self._cached_results.append(item)
+            yield item

tests/core/generator-utils/test_cached_iteratable.py

@@ -0,0 +1,25 @@
+from eth._utils.generator import CachedIterable
+
+from eth_utils.toolz import (
+    first,
+    nth,
+)
+import itertools
+
+
+def test_cached_generator():
+    use_once = itertools.count()
+    repeated_use = CachedIterable(use_once)
+
+    for find_val in [1, 0, 10, 5]:
+        assert find_val == nth(find_val, repeated_use)
+
+
+def test_laziness():
+    def crash_after_first_val():
+        yield 1
+        raise Exception("oops, iterated past first value")
+
+    repeated_use = CachedIterable(crash_after_first_val())
+    assert first(repeated_use) == 1
+    assert first(repeated_use) == 1