From fda40cb64041aacdb776e0b1f4f4a635bdc9d70b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Barrois?= Date: Sat, 15 Jun 2013 00:17:41 +0200 Subject: Add Iterator.reset() (Closes #63). --- factory/utils.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) (limited to 'factory/utils.py') 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) -- cgit v1.2.3