summaryrefslogtreecommitdiff
path: root/factory
diff options
context:
space:
mode:
Diffstat (limited to 'factory')
-rw-r--r--factory/__init__.py4
-rw-r--r--factory/declarations.py39
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)