diff options
Diffstat (limited to 'factory')
-rw-r--r-- | factory/declarations.py | 11 | ||||
-rw-r--r-- | factory/utils.py | 23 |
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) |