diff --git a/docs/api.rst b/docs/api.rst index 2b4902e8..553b7cf4 100644 --- a/docs/api.rst +++ b/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 ---- diff --git a/more_itertools/more.py b/more_itertools/more.py index fced0e36..4177ba4d 100644 --- a/more_itertools/more.py +++ b/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) diff --git a/more_itertools/tests/test_more.py b/more_itertools/tests/test_more.py index 04e3b751..e857e7cd 100644 --- a/more_itertools/tests/test_more.py +++ b/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)