Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Run 2to3 on IPython notebooks
#!/usr/bin/env python3
"""
To run: python3 nb2to3.py notebook-or-directory
"""
# Authors: Thomas Kluyver, Fernando Perez
# See: https://gist.github.com/takluyver/c8839593c615bb2f6e80
import argparse
import pathlib
from nbformat import read, write
import lib2to3
from lib2to3.refactor import RefactoringTool, get_fixers_from_package
def refactor_notebook_inplace(rt, path):
def refactor_cell(src):
#print('\n***SRC***\n', src)
try:
tree = rt.refactor_string(src+'\n', str(path) + '/cell-%d' % i)
except (lib2to3.pgen2.parse.ParseError,
lib2to3.pgen2.tokenize.TokenError):
return src
else:
return str(tree)[:-1]
print("Refactoring:", path)
nb = read(str(path), as_version=4)
# Run 2to3 on code
for i, cell in enumerate(nb.cells, start=1):
if cell.cell_type == 'code':
if cell.execution_count in (' ', '*'):
cell.execution_count = None
if cell.source.startswith('%%'):
# For cell magics, try to refactor the body, in case it's
# valid python
head, source = cell.source.split('\n', 1)
cell.source = head + '\n' + refactor_cell(source)
else:
cell.source = refactor_cell(cell.source)
# Update notebook metadata
nb.metadata.kernelspec = {
'display_name': 'Python 3',
'name': 'python3',
'language': 'python',
}
if 'language_info' in nb.metadata:
nb.metadata.language_info.codemirror_mode = {
'name': 'ipython',
'version': 3,
}
nb.metadata.language_info.pygments_lexer = 'ipython3'
nb.metadata.language_info.pop('version', None)
write(nb, str(path))
def main(argv=None):
ap = argparse.ArgumentParser()
ap.add_argument('path', type=pathlib.Path,
help="Notebook or directory containing notebooks")
options = ap.parse_args(argv)
avail_fixes = set(get_fixers_from_package('lib2to3.fixes'))
rt = RefactoringTool(avail_fixes)
if options.path.is_dir():
for nb_path in options.path.rglob('*.ipynb'):
refactor_notebook_inplace(rt, nb_path)
else:
refactor_notebook_inplace(rt, options.path)
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.