-
Notifications
You must be signed in to change notification settings - Fork 14
/
nbrun.py
87 lines (73 loc) · 3.04 KB
/
nbrun.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
# Copyright (c) 2015-2016 Antonino Ingargiola <tritemio@gmail.com>
# License: MIT
"""
Script to execute notebooks in a folder, unless their name
starts with '_' or ends with '-out'.
The executed notebooks are saved with a '-out' suffix.
Usage:
nbrun.py notebook_folder ouput_folder
"""
import os
import time
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
def run_notebook(notebook_name, nb_suffix='-out', out_path='.', timeout=3600,
**execute_kwargs):
"""Runs a notebook and saves the output in the same notebook.
Arguments:
notebook_name (string): name of the notebook to be executed.
nb_suffix (string): suffix to append to the file name of the executed
notebook.
timeout (int): timeout in seconds after which the execution is aborted.
execute_kwargs (dict): additional arguments passed to
`ExecutePreprocessor`.
out_path (string): folder where to save the output notebook.
"""
timestamp_cell = "**Executed:** %s\n\n**Duration:** %d seconds."
if str(notebook_name).endswith('.ipynb'):
notebook_name = str(notebook_name)[:-len('.ipynb')]
nb_name_input = notebook_name + '.ipynb'
nb_name_output = notebook_name + '%s.ipynb' % nb_suffix
nb_name_output = os.path.join(out_path, nb_name_output)
print('- Executing: ', nb_name_input)
execute_kwargs_ = dict(kernel_name = 'python%d' % sys.version_info[0])
if execute_kwargs is not None:
execute_kwargs_.update(execute_kwargs)
ep = ExecutePreprocessor(timeout=timeout, **execute_kwargs_)
nb = nbformat.read(nb_name_input, as_version=4)
start_time = time.time()
try:
# Execute the notebook
ep.preprocess(nb, {'metadata': {'path': './'}})
except:
# Execution failed, print a message then raise.
msg = 'Error executing the notebook "%s".\n\n' % notebook_name
msg += 'See notebook "%s" for the traceback.' % nb_name_output
print(msg)
raise
else:
# On successful execution, add timestamping cell
duration = time.time() - start_time
timestamp_cell = timestamp_cell % (time.ctime(start_time), duration)
nb['cells'].insert(0, nbformat.v4.new_markdown_cell(timestamp_cell))
finally:
# Save the notebook even when it raises an error
nbformat.write(nb, nb_name_output)
print('* Output: ', nb_name_output)
if __name__ == '__main__':
from pathlib import Path
import sys
path = '.'
if len(sys.argv) > 1:
path = sys.argv[1]
assert os.path.isdir(path), 'Folder "%s" not found.' % path
out_path = path
if len(sys.argv) > 2:
out_path = sys.argv[2]
assert os.path.isdir(out_path), 'Folder "%s" not found.' % out_path
print('Executing notebooks in "%s"... ' % os.path.abspath(path))
pathlist = list(Path(path).glob('*.ipynb'))
for nbpath in pathlist:
if not (nbpath.stem.endswith('-out') or nbpath.stem.startswith('_')):
print()
run_notebook(nbpath, out_path=out_path)