Add chunked_filter for batch predicates

boltons/iterutils.py

@@ -358,6 +358,45 @@ def postprocess(chk): return bytes(chk)
     return
 
 
+def chunked_filter(src, size, key=None):
+    """Yield items from *src* selected by a *key* function that operates on
+    chunks of up to *size* items. This is useful when a predicate can look
+    up several items in one database query or API request, avoiding a
+    separate request for every item.
+
+    The *key* function receives each chunk from :func:`chunked_iter` and
+    must return an iterable with one truth value per item, in the same
+    order. Items with a true value are yielded in their original order.
+    If *key* is ``None``, the items themselves are tested, as with
+    :func:`filter`.
+
+    For example, a batch lookup can identify records already processed
+    by another service:
+
+    >>> processed_ids = {2, 5}
+    >>> def is_new(ids):
+    ...     return [item_id not in processed_ids for item_id in ids]
+    >>> list(chunked_filter(range(1, 8), 3, key=is_new))
+    [1, 3, 4, 6, 7]
+
+    The source is consumed one chunk at a time, only as output is
+    requested. The final chunk may have fewer than *size* items. A
+    :exc:`ValueError` is raised if *key* returns the wrong number of
+    values for a chunk.
+    """
+    if key is not None and not callable(key):
+        raise TypeError('expected a callable key or None, not %r' % (key,))
+
+    for chunk in chunked_iter(src, size):
+        allowed = chunk if key is None else list(key(chunk))
+        if len(allowed) != len(chunk):
+            raise ValueError('chunked_filter expected key to return %d values, got %d'
+                             % (len(chunk), len(allowed)))
+        for item, allow in zip(chunk, allowed):
+            if allow:
+                yield item
+
+
 def chunk_ranges(input_size, chunk_size, input_offset=0, overlap_size=0, align=False):
     """Generates *chunk_size*-sized chunk ranges for an input with length *input_size*.
     Optionally, a start of the input can be set via *input_offset*, and

docs/iterutils.rst

@@ -18,6 +18,7 @@ present in the standard library.
 
 .. autofunction:: chunked
 .. autofunction:: chunked_iter
+.. autofunction:: chunked_filter
 .. autofunction:: chunk_ranges
 .. autofunction:: pairwise
 .. autofunction:: pairwise_iter

tests/test_iterutils.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from boltons import iterutils
 from boltons.dictutils import OMD
 from boltons.iterutils import (first,
                                split,
@@ -25,6 +26,93 @@
 is_meaning_of_life = lambda x: x == 42
 
 
+class TestChunkedFilter:
+    def test_batch_predicate(self):
+        records = [{'id': value} for value in [1, 2, 3, 2, 4, 5, 6]]
+        requested = []
+        existing = {2, 5}
+
+        def key(chunk):
+            requested.append([record['id'] for record in chunk])
+            return (record['id'] not in existing for record in chunk)
+
+        result = list(iterutils.chunked_filter(records, 3, key=key))
+
+        assert requested == [[1, 2, 3], [2, 4, 5], [6]]
+        assert len(result) == 4
+        assert all(actual is records[index]
+                   for actual, index in zip(result, [0, 2, 4, 6]))
+
+    def test_lazy_consumption(self):
+        consumed = []
+        checked = []
+
+        def source():
+            for item in range(7):
+                consumed.append(item)
+                yield item
+
+        def key(chunk):
+            checked.append(chunk)
+            return [item % 2 == 0 for item in chunk]
+
+        result = iterutils.chunked_filter(source(), 3, key=key)
+        assert consumed == checked == []
+        assert next(result) == 0
+        assert consumed == [0, 1, 2]
+        assert checked == [[0, 1, 2]]
+        assert next(result) == 2
+        assert consumed == [0, 1, 2]
+        assert next(result) == 4
+        assert consumed == [0, 1, 2, 3, 4, 5]
+        assert checked == [[0, 1, 2], [3, 4, 5]]
+        assert list(result) == [6]
+        assert checked == [[0, 1, 2], [3, 4, 5], [6]]
+
+    @pytest.mark.parametrize('src, expected', [
+        ([0, 1, None, 2, '', 3], [1, 2, 3]),
+        ('abc', ['a', 'b', 'c']),
+        (b'\x00ab', [97, 98]),
+    ])
+    def test_default_key(self, src, expected):
+        assert list(iterutils.chunked_filter(src, 2)) == expected
+        assert list(iterutils.chunked_filter(iter(src), 2, key=None)) == expected
+
+    def test_predicate_truth_values(self):
+        def key(chunk):
+            return [0, None, 'yes', [1]]
+
+        assert list(iterutils.chunked_filter(range(4), 4, key)) == [2, 3]
+
+    def test_empty(self):
+        def key(chunk):
+            pytest.fail('key must not be called for an empty source')
+
+        assert list(iterutils.chunked_filter(iter([]), 3, key)) == []
+
+    @pytest.mark.parametrize('mask', [[], [True], [True, False, True]])
+    def test_wrong_predicate_length(self, mask):
+        with pytest.raises(ValueError, match='expected key to return 2 values'):
+            list(iterutils.chunked_filter([1, 2], 2, lambda chunk: iter(mask)))
+
+    @pytest.mark.parametrize('key', [False, 1, 'key'])
+    def test_noncallable_key(self, key):
+        with pytest.raises(TypeError, match='expected a callable key'):
+            list(iterutils.chunked_filter([1, 2], 2, key))
+
+    @pytest.mark.parametrize('size', [0, -1])
+    def test_invalid_size(self, size):
+        with pytest.raises(ValueError, match='positive integer'):
+            list(iterutils.chunked_filter([1], size))
+
+    def test_key_exception(self):
+        def key(chunk):
+            raise RuntimeError('batch lookup failed')
+
+        with pytest.raises(RuntimeError, match='batch lookup failed'):
+            list(iterutils.chunked_filter([1], 2, key))
+
+
 class TestSplit:
     def test_maxsplit_zero_returns_unsplit_values(self):
         values = [1, None, 2]