Merge pull request #113 from bartvm/constant_iterator

Re-introduce constant iterator

blocks/datasets/schemes.py

@@ -1,4 +1,3 @@
-import itertools
 from abc import ABCMeta, abstractmethod
 
 import six
@@ -70,10 +69,26 @@ def __init__(self, batch_size, times=None):
         self.times = times
 
     def get_request_iterator(self):
-        if self.times is None:
-            return itertools.repeat(self.batch_size)
-        else:
-            return itertools.repeat(self.batch_size, self.times)
+        return ConstantIterator(self.batch_size, self.times)
+
+
+class ConstantIterator(six.Iterator):
+    def __init__(self, batch_size, times=None):
+        self.batch_size = batch_size
+        self.times = times
+        if times is not None:
+            self.current = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.times is not None:
+            if self.current == self.times:
+                raise StopIteration
+            else:
+                self.current += 1
+        return self.batch_size
 
 
 class SequentialScheme(BatchScheme):