Add locate() (#140)

docs/api.rst

@@ -113,7 +113,7 @@ These tools combine multiple iterables.
 Summarizing
 ===========
 
-These tools return summarized or aggregate data from an iterable.
+These tools return summarized or aggregated data from an iterable.
 
 ----
 
@@ -123,6 +123,7 @@ These tools return summarized or aggregate data from an iterable.
 .. autofunction:: first(iterable[, default])
 .. autofunction:: one
 .. autofunction:: unique_to_each
+.. autofunction:: locate
 
 ----
 

more_itertools/more.py

@@ -3,7 +3,16 @@
 from collections import Counter, defaultdict, deque
 from functools import partial, wraps
 from heapq import merge
-from itertools import chain, count, groupby, islice, repeat, takewhile, tee
+from itertools import (
+    chain,
+    compress,
+    count,
+    groupby,
+    islice,
+    repeat,
+    takewhile,
+    tee
+)
 from operator import itemgetter, lt, gt
 from sys import version_info
 
@@ -31,6 +40,7 @@
     'interleave',
     'intersperse',
     'iterate',
+    'locate',
     'numeric_range',
     'one',
     'padded',
@@ -1292,3 +1302,22 @@ def count_cycle(iterable, n=None):
         return iter(())
     counter = count() if n is None else range(n)
     return ((i, item) for i in counter for item in iterable)
+
+
+def locate(iterable, pred=bool):
+    """Yield the index of each item in *iterable* for which *pred* returns
+    ``True``.
+
+    *pred* defaults to :func:`bool`, which will select truthy items:
+
+        >>> list(locate([0, 1, 1, 0, 1, 0, 0]))
+        [1, 2, 4]
+
+    Set *pred* to a custom function to, e.g., find the indexes for a particular
+    item:
+
+        >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b'))
+        [1, 3]
+
+    """
+    return compress(count(), map(pred, iterable))

more_itertools/tests/test_more.py

@@ -1212,3 +1212,23 @@ def test_empty(self):
 
     def test_negative(self):
         self.assertEqual(list(count_cycle('abc', -3)), [])
+
+
+class LocateTests(TestCase):
+    def test_default_pred(self):
+        iterable = [0, 1, 1, 0, 1, 0, 0]
+        actual = list(locate(iterable))
+        expected = [1, 2, 4]
+        self.assertEqual(actual, expected)
+
+    def test_no_matches(self):
+        iterable = [0, 0, 0]
+        actual = list(locate(iterable))
+        expected = []
+
+    def test_custom_pred(self):
+        iterable = ['0', 1, 1, '0', 1, '0', '0']
+        pred = lambda x: x == '0'
+        actual = list(locate(iterable, pred))
+        expected = [0, 3, 5, 6]
+        self.assertEqual(actual, expected)