zonca · October 3, 2013 00:11 · Aug 5, 2013 · Mar 20, 2013
diff --git a/ipynb_run_save.py b/ipynb_run_save.py
@@ -15,13 +15,19 @@
 
 from Queue import Empty
 
-from IPython.kernel.blockingkernelmanager import BlockingKernelManager
+try:
+    # 1.0
+    from IPython.kernel import KernelManager
+except ImportError:
+    # 0.13
+    from IPython.zmq.blockingkernelmanager import BlockingKernelManager
+
 from IPython.nbformat.current import read, write, NotebookNode
 
 
-def run_cell(km, cell):
-    shell = km.shell_channel
-    iopub = km.iopub_channel
+def run_cell(kc, cell):
+    shell = kc.shell_channel
+    iopub = kc.iopub_channel
     # print "\n\ntesting:"
     # print cell.input
     shell.execute(cell.input)
@@ -68,16 +74,21 @@ def run_cell(km, cell):
 
 
 def test_notebook(nb):
-    km = BlockingKernelManager()
+    km = KernelManager()
     km.start_kernel(extra_arguments=['--pylab=inline'], stderr=open(os.devnull, 'w'))
-    km.start_channels()
+    try:
+        kc = km.client()
+    except AttributeError:
+        # 0.13
+        kc = km
+    kc.start_channels()
     # run %pylab inline, because some notebooks assume this
     # even though they shouldn't
-    km.shell_channel.execute("pass")
-    km.shell_channel.get_msg()
+    kc.shell_channel.execute("pass")
+    kc.shell_channel.get_msg()
     while True:
         try:
-            km.iopub_channel.get_msg(timeout=1)
+            kc.iopub_channel.get_msg(timeout=1)
         except Empty:
             break
 
@@ -89,7 +100,7 @@ def test_notebook(nb):
                 continue
             cells += 1
             try:
-                outs = run_cell(km, cell)
+                outs = run_cell(kc, cell)
             except Exception as e:
                 print "failed to run cell:", repr(e)
                 print cell.input

diff --git a/ipynb_run_save.py b/ipynb_run_save.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+"""
+Script for running and save notebooks from command line.
+
+How to use: `ipynb_run_save.py foo.ipynb
+
+Some tweaks over ipydoctest.py from minrk
+
+by @damianavila
+"""
+
+import io
+import os
+import sys
+
+from Queue import Empty
+
+from IPython.kernel.blockingkernelmanager import BlockingKernelManager
+from IPython.nbformat.current import read, write, NotebookNode
+
+
+def run_cell(km, cell):
+    shell = km.shell_channel
+    iopub = km.iopub_channel
+    # print "\n\ntesting:"
+    # 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'):
+            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 = BlockingKernelManager()
+    km.start_kernel(extra_arguments=['--pylab=inline'], stderr=open(os.devnull, 'w'))
+    km.start_channels()
+    # run %pylab inline, because some notebooks assume this
+    # even though they shouldn't
+    km.shell_channel.execute("pass")
+    km.shell_channel.get_msg()
+    while True:
+        try:
+            km.iopub_channel.get_msg(timeout=1)
+        except Empty:
+            break
+
+    errors = 0
+    cells = 0
+    for ws in nb.worksheets:
+        for cell in ws.cells:
+            if cell.cell_type != 'code':
+                continue
+            cells += 1
+            try:
+                outs = run_cell(km, cell)
+            except Exception as e:
+                print "failed to run cell:", repr(e)
+                print cell.input
+                errors += 1
+                continue
+            cell.outputs = outs
+
+    if errors:
+        print "    %3i cells failed to complete" % errors
+    if cells:
+        print "%i code cells from notebook %s" % (cells, nb.metadata.name)
+    km.shutdown_kernel()
+    del km
+
+if __name__ == '__main__':
+    for ipynb in sys.argv[1:]:
+        print "running %s" % ipynb
+        with io.open(ipynb, encoding='utf8') as f:
+            nb = read(f, 'json')
+        test_notebook(nb)
+        base, ext = os.path.splitext(ipynb)
+        new_ipynb = "%s_run_saved%s" % (base, ext)
+        with io.open(new_ipynb, 'w', encoding='utf8') as f:
+            write(nb, f, 'json')
+        print "wrote %s" % new_ipynb