Automatic Jupyter Notebook Python cell formatting. Runs YAPF on individual notebook cells using yapf_api.

README.adoc

yapf_nbformat

Runs YAPF on individual notebook cells using yapf_api.

• To skip a cell formatting (for example with jupyter-custom syntax, add # noqa to the first line of the cell.

• Remove newlines at end of cells.

• Uses the style_format file .style.yapf.

• Has a dry run option for checking proper formatting (f.e. in a Git hook).

yapf_nbformat.py

import nbformat
from yapf.yapflib.yapf_api import FormatCode

style_file = '.style.yapf'


def format_nb(notebook_filename, dry_run=False):
    print('Formatting {}...'.format(notebook_filename), end='')
    with open(notebook_filename, 'r') as f:
        notebook = nbformat.read(f, as_version=nbformat.NO_CONVERT)
    nbformat.validate(notebook)

    changed = False

    for cell in notebook.cells:
        if cell['cell_type'] != 'code':
            continue

        src = cell['source']
        lines = src.split('\n')
        if len(lines) <= 0 or '# noqa' in lines[0]:
            continue

        formatted_src, did_change = FormatCode(src, style_config=style_file)
        if formatted_src.endswith('\n'):
            formatted_src = formatted_src[:-1]  # remove last newline
            did_change = True
        if did_change:
            cell['source'] = formatted_src
            changed = True

    if changed:
        if dry_run:
            print(' (reformatted)')
        else:
            with open(notebook_filename, 'w') as f:
                nbformat.write(notebook, f, version=nbformat.NO_CONVERT)
            print()
    else:
        print()


def main(notebook_filenames, dry_run=False):
    for fn in notebook_filenames:
        format_nb(fn, dry_run=dry_run)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('file_to_format', nargs='+', help='The jupyter notebook file to format in place.')
    parser.add_argument(
        '--dry_run', action='store_true', help='Whether to just print if a notebook would be reformatted.')
    args = parser.parse_args()

    main(args.file_to_format, dry_run=args.dry_run)