jupyter-run

jupyter-run.py

#!/usr/bin/env python
import os
import sys
import nbformat
import nbformat.sign
from nbconvert.preprocessors import ExecutePreprocessor
from nbconvert.preprocessors.execute import CellExecutionError
import warnings

STDIO_CODE = u'''
class Tee:
    def __init__(self, *files):
        self._files = files

    def __del__(self):
        pass

    def write(self, string):
        for file in self._files:
            file.write(string)

    def flush(self):
        for file in self._files:
            file.flush()

sys.stdin = Tee(sys.__stdin__, sys.stdin)
sys.stdout = Tee(sys.__stdout__, sys.stdout)
sys.stderr = Tee(sys.__stderr__, sys.stderr)
del sys, os, Tee
'''

def runnb(nb_path, *args, kernel_name=None):
    with open(nb_path, 'r') as nb_file:
        nb = nbformat.read(nb_file, nbformat.NO_CONVERT)
    notary = nbformat.sign.NotebookNotary()
    trusted = notary.check_signature(nb)
    ep = ExecutePreprocessor(timeout=-1, allow_errors=False, kernel_name=kernel_name)
    try:
        stdio_code = "import sys, os\nsys.argv = {!r}\n__file__=sys.argv[0]".format([nb_path]+list(args))
        cell = nbformat.v4.new_code_cell(source=stdio_code+STDIO_CODE)
        nb.cells.insert(0, cell)
        _ = ep.preprocess(nb, {'metadata':{'path':os.getcwd()}})
    except CellExecutionError:
        print('Error executing the notebook "%s".' % nb_path, file=sys.stderr)
        raise
    finally:
        nb.cells.pop(0)
        for cell in nb.cells:
            if cell.cell_type == 'code' and cell.outputs:
                cell.metadata.collapsed = False
                cell.metadata.autoscroll = 'auto'
        if trusted:
            notary.sign(nb)
        else:
            notary.unsign(nb)
        with open(nb_path, mode='wt') as f:
            nbformat.write(nb, f)

if __name__ == '__main__':
    warnings.filterwarnings('error', '.*trusted.*', DeprecationWarning, __name__)
    import argparse
    parser = argparse.ArgumentParser(description='Execute a Jupyter notebook.')
    parser.add_argument('notebook', help='Path to the notebook to execute.')
    parser.add_argument('--kernel', help='Name of the kernel to use.', default=None)
    parser.add_argument('args', nargs=argparse.REMAINDER, help='Arguments to pass to the notebook.')
    args = parser.parse_args()
    runnb(args.notebook, *args.args, kernel_name=args.kernel)