diff options
Diffstat (limited to 'factory')
-rw-r--r-- | factory/__init__.py | 4 | ||||
-rw-r--r-- | factory/declarations.py | 39 |
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) |