Add safe generator

dimagi/utils/couch/database.py

@@ -7,6 +7,8 @@
 from requests.exceptions import RequestException
 from time import sleep
 
+from dimagi.utils.safe_iterator import safe_generator
+
 
 class DocTypeMismatchException(Exception):
     pass
@@ -58,12 +60,15 @@ def get_view_names(database):
             views.append("%s/%s" % (doc.name, view_name))
     return views
 
+
+@safe_generator
 def iter_docs(database, ids, chunksize=100, **query_params):
     for doc_ids in chunked(ids, chunksize):
         for doc in get_docs(database, keys=doc_ids, **query_params):
             yield doc
 
 
+@safe_generator
 def iter_docs_with_retry(database, ids, chunksize=100, max_attempts=5, **query_params):
     """
     A version of iter_docs that retries fetching documents if the connection

dimagi/utils/safe_iterator.py

@@ -0,0 +1,56 @@
+from functools import wraps
+
+
+class IteratorAlreadyConsumedException(Exception):
+    pass
+
+
+class SafeIterator(object):
+    """
+    This iterator-like object wraps an iterator.
+    The SafeIterator will raise a IteratorAlreadyConsumedException if a user attempts to iterate through it a
+    second time.
+    This is useful if the wrapped iterator is a one time user generator, which would simply raise StopIteration
+    exception again (thereby, appearing "empty" to a user who didn't realize that the generator had already been
+    consumed).
+
+    Note that this class cannot be used interchangeably with an iterator, because iterators must raise
+    StopIteration exceptions on subsequent iterations. However, in practice, I don't expect this to be an issue.
+    https://docs.python.org/2/library/stdtypes.html#iterator-types
+    """
+    def __init__(self, generator):
+        """
+        Create a new SafeIterator by wrapping a normal generator (or iterator)
+        :param generator: the generator to be wrapped
+        """
+        self._generator = generator
+        self._has_been_consumed = False
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self._has_been_consumed:
+            raise IteratorAlreadyConsumedException
+        try:
+            return self._generator.next()
+        except StopIteration as e:
+            self._has_been_consumed = True
+            raise
+
+    def __next__(self):
+        # Python 3 compatibility
+        return self.next()
+
+
+def safe_generator(fn):
+    """
+    This decorator function wraps the return value of the given function in a SafeIterator.
+    You should only use this decorator on functions that return generators.
+    The SafeIterator will render the generator "safe" by raising a IteratorAlreadyConsumedException if it is
+    iterated a second time.
+    """
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        return SafeIterator(fn(*args, **kwargs))
+    return wrapper

dimagi/utils/tests/sage_iterator.py

@@ -0,0 +1,24 @@
+from django.test import SimpleTestCase
+
+from dimagi.utils.safe_iterator import SafeIterator, IteratorAlreadyConsumedException, safe_generator
+
+
+class SafeIteratorTest(SimpleTestCase):
+
+    def test_exception(self):
+        generator = (x for x in range(3))
+        safe_generator = SafeIterator(generator)
+        self.assertEqual(list(safe_generator), [0, 1, 2])
+        with self.assertRaises(IteratorAlreadyConsumedException):
+            for i in safe_generator:
+                pass
+
+    def test_decorator(self):
+        @safe_generator
+        def func():
+            return (x for x in range(3))
+
+        generator = func()
+        self.assertEqual(list(generator), [0, 1, 2])
+        with self.assertRaises(IteratorAlreadyConsumedException):
+            list(generator)