iter_batches.py

def iter_batches(iterable, batch_size):
  """Iterates over the given iterable in batches.

  Args:
      iterable: an iterable
      batch_size: the desired batch size, or None to return the contents in
          a single batch

  Returns:
      a generator that emits tuples of elements of the requested batch size
      from the input
  """
  it = iter(iterable)
  while True:
    chunk = tuple(itertools.islice(it, batch_size))
    if not chunk:
      return
    yield chunk