batch_iterator.py

class BatchIterator(object):
    """
    Return batches of items of batch size or less from the given source.
    
    Each batch is another iterator.
    """
    def __init__(self, source, batch_size=900):
        self.source_iter = iter(source)
        self.batch_size = batch_size
        self.stop_iteration = False

    def __iter__(self):
        return self

    def __next__(self):
        if self.stop_iteration:
            raise StopIteration
        def batch():
            for item in self.source_iter:
                self.index += 1
                yield item
                if self.index % self.batch_size == 0 and self.index != 0:
                    break
            else:
                self.stop_iteration = True
        return batch