Iterate by batches

batch.py

def batch(iterable, size, drop_remainder=False):
    if size < 1:
        raise ValueError("size must be greater than zero.")
        
    iterator = iter(iterable)
    while True:
        try:
            items = []
            for _ in range(size):
                items.append(next(iterator))
        except StopIteration:
            if len(items)==0 or drop_remainder:
                return None
        yield items