-
-
Save ankostis/8a6b3a9d2da18267577cd9b1fd626194 to your computer and use it in GitHub Desktop.
simple example script for running and testing notebooks in py3
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
Simple example script for running and testing notebooks. | |
Usage: `ipnbdoctest.py foo.ipynb [bar.ipynb [...]]` | |
Each cell is submitted to the kernel, and the outputs are compared with those stored in the notebook. | |
- Found in this SO: http://stackoverflow.com/a/20483736/548792 | |
- Copied from https://gist.github.com/fdeheeger/e79fac57e80f8aaa63cc | |
- Attempt to make it run under py3.4+. | |
License: Public Domain, but credit is nice (Min RK). | |
""" | |
import os,sys,time | |
import base64 | |
import re | |
import os | |
import os.path as osp | |
from collections import defaultdict | |
from queue import Queue, Empty | |
try: | |
from IPython.kernel import KernelManager | |
except ImportError: | |
from IPython.zmq.blockingkernelmanager import BlockingKernelManager as KernelManager | |
from nbformat.current import reads, NotebookNode | |
mydir = osp.dirname(__file__) | |
def compare_png(a64, b64): | |
"""compare two b64 PNGs (incomplete)""" | |
try: | |
import Image | |
except ImportError: | |
pass | |
adata = base64.decodestring(a64) | |
bdata = base64.decodestring(b64) | |
return adata == bdata | |
def sanitize(s): | |
"""sanitize a string for comparison. | |
fix universal newlines, strip trailing newlines, and normalize likely random values (memory addresses and UUIDs) | |
""" | |
if not isinstance(s, str): | |
return s | |
# normalize newline: | |
s = s.replace('\r\n', '\n') | |
# ignore trailing newlines (but not space) | |
s = s.rstrip('\n') | |
# normalize hex addresses: | |
s = re.sub(r'0x[a-f0-9]+', '0xFFFFFFFF', s) | |
# normalize UUIDs: | |
s = re.sub(r'[a-f0-9]{8}(\-[a-f0-9]{4}){3}\-[a-f0-9]{12}', 'U-U-I-D', s) | |
return s | |
def consolidate_outputs(outputs): | |
"""consolidate outputs into a summary dict (incomplete)""" | |
data = defaultdict(list) | |
data['stdout'] = '' | |
data['stderr'] = '' | |
for out in outputs: | |
if out.type == 'stream': | |
data[out.stream] += out.text | |
elif out.type == 'pyerr': | |
data['pyerr'] = dict(ename=out.ename, evalue=out.evalue) | |
else: | |
for key in ('png', 'svg', 'latex', 'html', 'javascript', 'text', 'jpeg',): | |
if key in out: | |
data[key].append(out[key]) | |
return data | |
def compare_outputs(test, ref, skip_compare=('png', 'traceback', 'latex', 'prompt_number')): | |
for key in ref: | |
if key not in test: | |
print("missing key: %s != %s" % (test.keys(), ref.keys())) | |
return False | |
elif key not in skip_compare and sanitize(test[key]) != sanitize(ref[key]): | |
print("mismatch %s:" % key) | |
print(test[key]) | |
print(' != ') | |
print(ref[key]) | |
return False | |
return True | |
def run_cell(shell, iopub, cell): | |
# print(cell.input) | |
shell.execute(cell.input) | |
# wait for finish, maximum 20s | |
shell.get_msg(timeout=20) | |
outs = [] | |
while True: | |
try: | |
msg = iopub.get_msg(timeout=0.2) | |
except Empty: | |
break | |
msg_type = msg['msg_type'] | |
if msg_type in ('status', 'pyin'): | |
continue | |
elif msg_type == 'clear_output': | |
outs = [] | |
continue | |
content = msg['content'] | |
# print(msg_type, content) | |
out = NotebookNode(output_type=msg_type) | |
if msg_type == 'stream': | |
out.stream = content['name'] | |
out.text = content['data'] | |
elif msg_type in ('display_data', 'pyout'): | |
out['metadata'] = content['metadata'] | |
for mime, data in content['data'].iteritems(): | |
attr = mime.split('/')[-1].lower() | |
# this gets most right, but fix svg+html, plain | |
attr = attr.replace('+xml', '').replace('plain', 'text') | |
setattr(out, attr, data) | |
if msg_type == 'pyout': | |
out.prompt_number = content['execution_count'] | |
elif msg_type == 'pyerr': | |
out.ename = content['ename'] | |
out.evalue = content['evalue'] | |
out.traceback = content['traceback'] | |
else: | |
print("unhandled iopub msg:", msg_type) | |
outs.append(out) | |
return outs | |
def _test_notebook(nb): | |
km = KernelManager() | |
km.start_kernel(extra_arguments=['--pylab=inline'], stderr=open(os.devnull, 'w')) | |
try: | |
kc = km.client() | |
kc.start_channels() | |
iopub = kc.iopub_channel | |
except AttributeError: | |
# IPython 0.13 | |
kc = km | |
kc.start_channels() | |
iopub = kc.sub_channel | |
shell = kc.shell_channel | |
# run %pylab inline, because some notebooks assume this | |
# even though they shouldn't | |
shell.execute("pass") | |
shell.get_msg() | |
while True: | |
try: | |
iopub.get_msg(timeout=1) | |
except Empty: | |
break | |
successes = 0 | |
failures = 0 | |
errors = 0 | |
for ws in nb.worksheets: | |
for cell in ws.cells: | |
if cell.cell_type != 'code': | |
continue | |
try: | |
outs = run_cell(shell, iopub, cell) | |
except Exception as e: | |
print("failed to run cell:", repr(e)) | |
print(cell.input) | |
errors += 1 | |
continue | |
failed = False | |
for out, ref in zip(outs, cell.outputs): | |
if not compare_outputs(out, ref): | |
failed = True | |
if failed: | |
failures += 1 | |
else: | |
successes += 1 | |
sys.stdout.write('.') | |
print("tested notebook %s" % nb.metadata.name) | |
print(" %3i cells successfully replicated" % successes) | |
if failures: | |
print(" %3i cells mismatched output" % failures) | |
if errors: | |
print(" %3i cells failed to complete" % errors) | |
kc.stop_channels() | |
km.shutdown_kernel() | |
del km | |
if __name__ == '__main__': | |
for ipynb in sys.argv[1:]: | |
print("testing %s" % ipynb) | |
with open(ipynb) as f: | |
nb = reads(f.read(), 'json') | |
test_notebook(nb) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment