Add item_index() and sub_index() itertools

docs/api.rst

@@ -200,6 +200,7 @@ Others
 .. autofunction:: numeric_range(start, stop, step)
 .. autofunction:: side_effect
 .. autofunction:: iterate
+.. autofunction:: sub_index
 
 ----
 

more_itertools/more.py

@@ -31,6 +31,7 @@
     'interleave',
     'intersperse',
     'iterate',
+    'item_index',
     'numeric_range',
     'one',
     'padded',
@@ -42,6 +43,7 @@
     'split_before',
     'spy',
     'stagger',
+    'sub_index',
     'unique_to_each',
     'windowed',
     'with_iter',
@@ -1292,3 +1294,62 @@ 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 item_index(iterable, item, start=0, end=None):
+    """An extension of :func:`list.index` for iterables that returns the index
+    of *item* in *iterable*. Raises ``ValueError`` if it is not present.
+
+        >>> item_index(['mercury', 'venus', 'earth', 'mars'], 'earth')
+        2
+
+    Specify *start* and/or *end* to restrict the search:
+
+        >>> iterable = ['a', 'b', 'c', 'd', 'a', 'b', 'c', 'd']
+        >>> item = 'b'
+        >>> item_index(iterable, item, start=3, end=6)
+        5
+
+    Unlike its list-based counterpart, this function does not accept negative
+    *start* or *stop* values.
+
+    To search for a sub-sequence, see :func:`sub_index`.
+
+    """
+    it = filter(
+        lambda x: x[1] == item,
+        enumerate(islice(iterable, start, end), start),
+    )
+    return first(it)[0]
+
+
+def sub_index(iterable, sub, start=0, end=None):
+    """An extension of :func:`str.index` for iterables that returns the index
+    of sub-sequence *sub* in *iterable*. Raises ``ValueError`` if it is not
+    present.
+
+        >>> iterable = ['george', 'lucille', 'gob', 'michael', 'lindsay']
+        >>> sub = ['michael', 'lindsay']
+        >>> sub_index(iterable, sub)
+        3
+
+    Specify *start* and/or *end* to restrict the search:
+
+        >>> iterable = ['a', 'b', 'c', 'd', 'a', 'b', 'c', 'd']
+        >>> sub = ['a', 'b']
+        >>> sub_index(iterable, sub, start=1, end=6)
+        4
+
+    Unlike its str-based counterpart, this function does not accept negative
+    *start* or *stop* values.
+
+    Note that *sub* more be an iterable. It will be exhausted. To search for a
+    single item instead of a sub-sequence, see :func:`item_index`.
+
+    """
+    sub = tuple(sub)
+    it = filter(
+        lambda x: x[1] == sub,
+        enumerate(windowed(islice(iterable, start, end), len(sub)), start)
+    )
+    return first(it)[0]

more_itertools/tests/test_more.py

@@ -1212,3 +1212,44 @@ def test_empty(self):
 
     def test_negative(self):
         self.assertEqual(list(count_cycle('abc', -3)), [])
+
+
+class ItemIndexTests(TestCase):
+    def test_basic(self):
+        for item, kwargs, expected in [
+            ('a', {}, 0),
+            ('a', {'start': 1}, 8),
+            ('D', {'start': 8, 'end': 15}, None),
+            ('E', {'end': 16}, None),
+        ]:
+            iterable = (x for x in 'abcdABCDabcdABCDE')
+            try:
+                actual = item_index(iterable, item, **kwargs)
+            except ValueError:
+                actual = None
+
+            self.assertEqual(actual, expected)
+
+    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 [
+            ('abc', {}, 0),
+            ('abc', {'start': 1}, 8),
+            ('DE', {'start': 8, 'end': 15}, None),
+            ('E', {'end': 16}, None),
+        ]:
+            iterable = (x for x in 'abcdABCDabcdABCDE')
+            try:
+                actual = sub_index(iterable, sub, **kwargs)
+            except ValueError:
+                actual = None
+
+            self.assertEqual(actual, expected)
+
+    def test_invalid_index(self):
+        with self.assertRaises(ValueError):
+            sub_index(['a', 'b', 'c', 'd'], ['a', 'b'], start=-1)