jupyter-cleanup.py

#!/usr/bin/env python

import nbformat as nbf
import click
from itertools import count

@click.command(help='Cleanup a jupyter notebook, putting `execution_counts` in order')
@click.option('-i', '--input', type=click.File('r'), default='-', required=True)
@click.option('-o', '--output', type=click.File('w'), default='-', required=True)
def cli(input, output):
  nb = nbf.read(input, as_version=4)
  execution_count = iter(count(start=1))
  for cell in nb.cells:
    if 'execution_count' in cell:
      cell['execution_count'] = next(execution_count)
      for cell_output in cell['outputs']:
        if 'execution_count' in cell_output:
          cell_output['execution_count'] = cell['execution_count']
  nbf.write(nb, output)

if __name__ == '__main__':
  cli()