Merge pull request #279 from erikrose/time-limit

Add time_limited()

docs/api.rst

@@ -231,6 +231,7 @@ Others
 .. autofunction:: difference(iterable, func=operator.sub)
 .. autofunction:: make_decorator
 .. autoclass:: SequenceView
+.. autofunction:: time_limited
 
 ----
 

more_itertools/more.py

@@ -1,4 +1,5 @@
 from collections import Counter, defaultdict, deque
+from collections.abc import Sequence
 from functools import partial, wraps
 from heapq import merge
 from itertools import (
@@ -17,7 +18,7 @@
 )
 from operator import itemgetter, lt, gt, sub
 from sys import maxsize, version_info
-from collections.abc import Sequence
+from time import monotonic
 
 from .recipes import consume, flatten, powerset, take
 
@@ -74,6 +75,7 @@
     'strip',
     'substrings',
     'substrings_indexes',
+    'time_limited',
     'unique_to_each',
     'unzip',
     'windowed',
@@ -2369,3 +2371,33 @@ def partitions(iterable):
     n = len(sequence)
     for i in powerset(range(1, n)):
         yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))]
+
+
+def time_limited(limit_seconds, iterable):
+    """
+    Yield items from *iterable* until *limit_seconds* have passed.
+
+    >>> from time import sleep
+    >>> def generator():
+    ...     yield 1
+    ...     yield 2
+    ...     sleep(0.2)
+    ...     yield 3
+    >>> iterable = generator()
+    >>> list(time_limited(0.1, iterable))
+    [1, 2]
+
+    Note that the time is checked before each item is yielded, and iteration
+    stops if  the time elapsed is greater than *limit_seconds*. If your time
+    limit is 1 second, but it takes 2 seconds to generate the first item from
+    the iterable, the function will run for 2 seconds and not yield anything.
+
+    """
+    if limit_seconds < 0:
+        raise ValueError('limit_seconds must be positive')
+
+    start_time = monotonic()
+    for item in iterable:
+        if monotonic() - start_time > limit_seconds:
+            break
+        yield item

more_itertools/tests/test_more.py

@@ -16,6 +16,7 @@
     repeat,
 )
 from operator import add, mul, itemgetter
+from time import sleep
 from unittest import TestCase
 
 import more_itertools as mi
@@ -2425,3 +2426,27 @@ def test_duplicates(self):
             [[1], [1], [1]],
         ]
         self.assertEqual(actual, expected)
+
+
+class TimeLimitedTests(TestCase):
+    def test_basic(self):
+        def generator():
+            yield 1
+            yield 2
+            sleep(0.2)
+            yield 3
+
+        iterable = generator()
+        actual = list(mi.time_limited(0.1, iterable))
+        expected = [1, 2]
+        self.assertEqual(actual, expected)
+
+    def test_zero_limit(self):
+        iterable = count()
+        actual = list(mi.time_limited(0, iterable))
+        expected = []
+        self.assertEqual(actual, expected)
+
+    def test_invalid_limit(self):
+        with self.assertRaises(ValueError):
+            list(mi.time_limited(-0.1, count()))