Cherry pick cells from Jupyter Notebook files.

jupytate.py

import sys
import argparse
import json


def save_nb(fileobj, cells, metadata, nbformat, nbformat_minor):
    json.dump({
        'cells': cells,
        'metadata': metadata,
        'nbformat': nbformat,
        'nbformat_minor': nbformat_minor,
    }, fileobj)


def excerpt_from_source(source):
    return ''.join(source).strip().replace('\n', ' ')[0:45]


def print_cell_excerpts(notebook):
    for cell_id, cell in enumerate(notebook['cells']):
        excerpt = excerpt_from_source(cell['source'])
        print('{:<5} {:<12} {}...'.format(cell_id,
                                          cell['cell_type'],
                                          excerpt))


def prompt_for_cells():
    print('Which cells do you want to cherry pick? E.g. 3:5.')
    answer = input('--> ')
    return slice(*map(int, answer.split(':')))


def cherry_pick(notebook):
    print_cell_excerpts(notebook)

    picked_cells = prompt_for_cells()
    save_nb(fileobj=args.output,
            cells=notebook['cells'][picked_cells],
            metadata=notebook['metadata'],
            nbformat=notebook['nbformat'],
            nbformat_minor=notebook['nbformat_minor'])


def main(argv):
    parser = argparse.ArgumentParser('Cherry pick cells from Jupyter Notebook files.')
    parser.add_argument('-f', '--file', type=argparse.FileType('r'))
    parser.add_argument('-o', '--output', type=argparse.FileType('w'))
    args = parser.parse_args(argv)

    notebook = json.load(args.file)
    cherry_pick(notebook)


if __name__ == '__main__':
    main(sys.argv[1:])