Add locate() itertool

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,9 @@ These tools return summarized or aggregate data from an iterable.
 .. autofunction:: first(iterable[, default])
 .. autofunction:: one
 .. autofunction:: unique_to_each
+.. autofunction:: item_index
+.. autofunction:: sub_index
+.. autofunction:: locate
 
 ----
 
@@ -200,7 +203,6 @@ Others
 .. autofunction:: numeric_range(start, stop, step)
 .. autofunction:: side_effect
 .. autofunction:: iterate
-.. autofunction:: sub_index
 
 ----
 

more_itertools/more.py

@@ -10,7 +10,7 @@
 from six import binary_type, string_types, text_type
 from six.moves import filter, map, range, zip, zip_longest
 
-from .recipes import flatten, take
+from .recipes import flatten, nth, take
 
 __all__ = [
     'adjacent',
@@ -32,6 +32,7 @@
     'intersperse',
     'iterate',
     'item_index',
+    'locate',
     'numeric_range',
     'one',
     'padded',
@@ -1353,3 +1354,30 @@ def sub_index(iterable, sub, start=0, end=None):
         enumerate(windowed(islice(iterable, start, end), len(sub)), start)
     )
     return first(it)[0]
+
+
+def locate(iterable, pred=bool, n=0):
+    """Return the index of the first item for which callable *pred* returns
+    ``True``, optionally skipping the first *n* such items.
+    Returns ``None`` if there are no such items.
+
+    *pred* defaults to ``bool``, which will select truthy items:
+
+        >>> iterable = [0, 1, 1, 0, 1, 0, 0]
+        >>> locate(iterable)
+        1
+        >>> locate(iterable, n=1)  # Skip the first match
+        2
+
+    This function can be used to determine the index `n`-th occurrence of an
+    item in an iterable, if we consider `n` to be  zero-based:
+
+        >>> iterable = '_a_aaaa'
+        >>> pred = lambda x: x == 'a'
+        >>> locate(iterable, pred, 0)  # The index of the 0th instance of 'a'
+        1
+        >>> locate(iterable, pred, 4)  # The index of the 4th instance of 'a'
+        6
+
+    """
+    return nth((i for i, item in enumerate(iterable) if pred(item)), n)

more_itertools/tests/test_more.py

@@ -1234,6 +1234,7 @@ def test_invalid_index(self):
         with self.assertRaises(ValueError):
             item_index('abcd', 'a', start=-1)
 
+
 class SubIndexTests(TestCase):
     def test_basic(self):
         for sub, kwargs, expected in [
@@ -1253,3 +1254,21 @@ def test_basic(self):
     def test_invalid_index(self):
         with self.assertRaises(ValueError):
             sub_index(['a', 'b', 'c', 'd'], ['a', 'b'], start=-1)
+
+
+class LocateTests(TestCase):
+    def test_default_pred(self):
+        iterable = [0, 1, 1, 0, 1, 0, 0]
+        self.assertEqual(locate(iterable), 1)
+        self.assertEqual(locate(iterable, n=1), 2)
+        self.assertEqual(locate(iterable, n=2), 4)
+        self.assertEqual(locate(iterable, n=3), None)
+
+    def test_custom_pred(self):
+        iterable = [0, 1, 1, 0, 1, 0, 0]
+        pred = lambda x: bool(x) == False
+        self.assertEqual(locate(iterable, pred), 0)
+        self.assertEqual(locate(iterable, pred, 1), 3)
+        self.assertEqual(locate(iterable, pred, 2), 5)
+        self.assertEqual(locate(iterable, pred, 3), 6)
+        self.assertEqual(locate(iterable, pred, 4), None)