Add Iterator.reset() (Closes #63).

2 files changed, 30 insertions, 4 deletions

diff --git a/factory/declarations.py b/factory/declarations.py
index f068c0d..552ddf2 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -163,17 +163,20 @@ class Iterator(OrderedDeclaration):
         self.getter = getter
 
         if cycle:
-            self.iterator = itertools.cycle(iterator)
-        else:
-            self.iterator = iter(iterator)
+            iterator = itertools.cycle(iterator)
+        self.iterator = utils.ResetableIterator(iterator)
 
     def evaluate(self, sequence, obj, create, extra=None, containers=()):
         logger.debug("Iterator: Fetching next value from %r", self.iterator)
-        value = next(self.iterator)
+        value = next(iter(self.iterator))
         if self.getter is None:
             return value
         return self.getter(value)
 
+    def reset(self):
+        """Reset the internal iterator."""
+        self.iterator.reset()
+
 
 class Sequence(OrderedDeclaration):
     """Specific OrderedDeclaration to use for 'sequenced' fields.
diff --git a/factory/utils.py b/factory/utils.py
index e1b265f..48c6eed 100644
--- a/factory/utils.py
+++ b/factory/utils.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+import collections
 
 #: String for splitting an attribute name into a
 #: (subfactory_name, subfactory_field) tuple.
@@ -101,3 +102,25 @@ def log_pprint(args=(), kwargs=None):
         [str(arg) for arg in args] +
         ['%s=%r' % item for item in kwargs.items()]
     )
+
+
+class ResetableIterator(object):
+    """An iterator wrapper that can be 'reset()' to its start."""
+    def __init__(self, iterator, **kwargs):
+        super(ResetableIterator, self).__init__(**kwargs)
+        self.iterator = iter(iterator)
+        self.past_elements = collections.deque()
+        self.next_elements = collections.deque()
+
+    def __iter__(self):
+        while True:
+            if self.next_elements:
+                yield self.next_elements.popleft()
+            else:
+                value = next(self.iterator)
+                self.past_elements.append(value)
+                yield value
+
+    def reset(self):
+        self.next_elements.clear()
+        self.next_elements.extend(self.past_elements)