Add the Iterator and InfiniteIterator attribute kinds.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

2 files changed, 43 insertions, 0 deletions

diff --git a/factory/__init__.py b/factory/__init__.py
index cf9cc3e..12f297e 100644
--- a/factory/__init__.py
+++ b/factory/__init__.py
@@ -52,6 +52,8 @@ from base import (
 
 from declarations import (
     LazyAttribute,
+    Iterator,
+    InfiniteIterator,
     Sequence,
     LazyAttributeSequence,
     SelfAttribute,
@@ -59,6 +61,8 @@ from declarations import (
     SubFactory,
 
     lazy_attribute,
+    iterator,
+    infinite_iterator,
     sequence,
     lazy_attribute_sequence,
     container_attribute,
diff --git a/factory/declarations.py b/factory/declarations.py
index 5fe427c..41d99a3 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -21,6 +21,9 @@
 # THE SOFTWARE.
 
 
+import itertools
+
+
 class OrderedDeclaration(object):
     """A factory declaration.
 
@@ -110,6 +113,34 @@ class SelfAttribute(OrderedDeclaration):
         return deepgetattr(obj, self.attribute_name, self.default)
 
 
+class Iterator(OrderedDeclaration):
+    """Fill this value using the values returned by an iterator.
+
+    Warning: the iterator should not end !
+
+    Attributes:
+        iterator (iterable): the iterator whose value should be used.
+    """
+
+    def __init__(self, iterator):
+        super(Iterator, self).__init__()
+        self.iterator = iter(iterator)
+
+    def evaluate(self, sequence, obj, containers=()):
+        return self.iterator.next()
+
+
+class InfiniteIterator(Iterator):
+    """Same as Iterator, but make the iterator infinite by cycling at the end.
+
+    Attributes:
+        iterator (iterable): the iterator, once made infinite.
+    """
+
+    def __init__(self, iterator):
+        return super(InfiniteIterator, self).__init__(itertools.cycle(iterator))
+
+
 class Sequence(OrderedDeclaration):
     """Specific OrderedDeclaration to use for 'sequenced' fields.
 
@@ -224,6 +255,14 @@ class SubFactory(OrderedDeclaration):
 def lazy_attribute(func):
     return LazyAttribute(func)
 
+def iterator(func):
+    """Turn a generator function into an iterator attribute."""
+    return Iterator(func())
+
+def infinite_iterator(func):
+    """Turn a generator function into an infinite iterator attribute."""
+    return InfiniteIterator(func())
+
 def sequence(func):
     return Sequence(func)