+def iterator_wrapper(func):
+ """Decorator handling iterate option"""
+ @wraps(func)
+ def decorated_function(instance, *args, **kwargs):
+ generator = func(instance, *args, **kwargs)
+ if not instance.iterate:
+ return list(generator)
+ instance._iterating = True
+ def iterator(gen):
+ try:
+ for item in gen:
+ yield item
+ finally:
+ instance._iterating = False
+ return iterator(generator)
+ return decorated_function
+
+