Format all code in a notebook with Black

blacken_notebook.py

# https://black.readthedocs.io/en/stable/
import json
import black

import re
magic_re = re.compile(r'^%', flags=re.MULTILINE)
ipymagic = '##IpyMagic##'

filename = 'notebook.ipynb'

def blacken_nb(filename, line_length=100, output_fname=None):
    if output_fname is None:
        output_fname = filename.replace('.ipynb', '-formatted.ipynb')
        assert filename != output_fname

    with open(filename) as f:
        nb = json.load(f)

    for cell in nb['cells']:
        if cell['cell_type'] != 'code':
            continue
        source = cell['source']
        if isinstance(source, list):
            source = ''.join(cell['source'])
        # Black doesn't know how to handle IPython magic, like '%matplotlib'.
        # So replace it with a specially formatted comment.
        source = magic_re.sub(ipymagic, source)

        formatted = black.format_str(source, line_length=line_length).rstrip('\n')

        # Replace the special comment with the magic.
        formatted = formatted.replace(ipymagic, '%')
        cell['source'] = formatted.splitlines(keepends=True)

    with open(output_fname, 'w') as f:
        json.dump(nb, f)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('filename')
    opts = parser.parse_args()
    blacken_nb(opts.filename)