-
Notifications
You must be signed in to change notification settings - Fork 240
/
docutils.py
111 lines (92 loc) · 3.74 KB
/
docutils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import json
import os
import sys
import argparse
from openmdao.utils.file_utils import files_iter
def reset_notebook(fname, dryrun=False):
"""
Empties the output fields and resets execution_count in all code cells in the given notebook.
Also removes any empty code cells. The specified notebook is overwritten.
Parameters
----------
fname : str
Name of the notebook file.
dryrun : bool
If True, don't actually update the file.
Returns
-------
bool
True if the file was updated or would have been updated if not a dry run.
"""
changed = False
with open(fname) as f:
dct = json.load(f)
newcells = []
for cell in dct['cells']:
if cell['cell_type'] == 'code':
if cell['source']: # cell is not empty
if cell['execution_count'] is not None or len(cell['outputs']) > 0:
cell['execution_count'] = None
cell['outputs'] = []
changed = True
newcells.append(cell)
else:
newcells.append(cell)
changed |= len(dct['cells']) != len(newcells)
dct['cells'] = newcells
if changed and not dryrun:
with open(fname, 'w') as f:
json.dump(dct, f, indent=1, ensure_ascii=False)
print(file=f) # avoid 'no newline at end of file' message
return changed
def reset_notebook_cmd():
"""
Run reset_notebook on notebook files.
"""
parser = argparse.ArgumentParser(description='Empty output cells, reset execution_count, and '
'remove empty cells of jupyter notebook(s).')
parser.add_argument('file', nargs='*', help='Jupyter notebook file(s).')
parser.add_argument('-r', '--recurse', action='store_true', dest='recurse',
help='Search through all directories at or below the current one for the '
'specified file(s). If no files are specified, reset all jupyter notebook '
'files found.')
parser.add_argument('-i', '--include', action='append', dest='includes',
default=[], help='If the --recurse option is active, this specifies a '
'local filename or glob pattern to match. This argument may be supplied '
'multiple times.')
parser.add_argument('-d', '--dryrun', action='store_true', dest='dryrun',
help="Report which notebooks would be updated but don't actually update "
"them.")
args = parser.parse_args()
if args.dryrun:
updatestr = 'Would have updated file'
else:
updatestr = 'Updated file'
if args.recurse:
if args.file:
print("When using the --recurse option, don't specify filenames. Use --include "
"instead.")
sys.exit(-1)
if not args.includes:
args.includes = ['*.ipynb']
for f in files_iter(file_includes=args.includes):
if not f.endswith('.ipynb'):
print(f"Ignoring {f} (not a notebook).")
continue
if reset_notebook(f, args.dryrun):
print(updatestr, f)
else:
if args.includes:
print("The --include option only works when also using --recurse.")
sys.exit(-1)
for f in sorted(args.file):
if os.path.isdir(f):
continue
if not f.endswith('.ipynb'):
print(f"Ignoring {f} (not a notebook).")
continue
if not os.path.isfile(f):
print(f"Can't find file '{f}'.")
sys.exit(-1)
if reset_notebook(f, args.dryrun):
print(updatestr, f)