From e30386ce4a01c7b03a341850239b0d0c39b2b80e Mon Sep 17 00:00:00 2001
From: Javi Merino <javi.merino@arm.com>
Date: Thu, 16 Apr 2015 19:24:50 +0100
Subject: [PATCH] Add ipython version 3 support for the generic ipython support

---
 wlauto/utils/ipython.py | 37 +++++++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/wlauto/utils/ipython.py b/wlauto/utils/ipython.py
index beacbb95..dfe84f98 100644
--- a/wlauto/utils/ipython.py
+++ b/wlauto/utils/ipython.py
@@ -28,20 +28,37 @@ except ImportError as import_error:
     else:
         import_error_str = import_error.message
 
+# The current code generates notebooks version 3
+NBFORMAT_VERSION = 3
+
 if IPython and (IPython.version_info[0] == 2):
     import IPython.kernel
     import IPython.nbformat.v3
 
     def read_notebook(notebook_in):
-        return IPython.nbformat.v3.reads_json(notebook_in)
+        return IPython.nbformat.v3.reads_json(notebook_in)  # pylint: disable=E1101
 
     def write_notebook(notebook, fout):
-        IPython.nbformat.v3.nbjson.JSONWriter().write(notebook, fout)
+        IPython.nbformat.v3.nbjson.JSONWriter().write(notebook, fout)  # pylint: disable=E1101
 
-    NotebookNode = IPython.nbformat.v3.NotebookNode
+    NotebookNode = IPython.nbformat.v3.NotebookNode  # pylint: disable=E1101
 
     IPYTHON_NBCONVERT = ['ipython', 'nbconvert', '--to=latex', '--post=PDF']
 
+elif IPython and (IPython.version_info[0] == 3):
+    import IPython.kernel
+    import IPython.nbformat
+
+    def read_notebook(notebook_in):
+        return IPython.nbformat.reads(notebook_in, NBFORMAT_VERSION)  # pylint: disable=E1101
+
+    def write_notebook(notebook, fout):
+        IPython.nbformat.write(notebook, fout)  # pylint: disable=E1101
+
+    NotebookNode = IPython.nbformat.NotebookNode  # pylint: disable=E1101
+
+    IPYTHON_NBCONVERT = ['ipython', 'nbconvert', '--to=pdf']
+
 elif IPython:
     # Unsupported IPython version
     IPython_ver_str = ".".join([str(n) for n in IPython.version_info])
@@ -51,12 +68,20 @@ elif IPython:
 def parse_valid_output(msg):
     """Parse a valid result from an execution of a cell in an ipython kernel"""
     msg_type = msg["msg_type"]
+    if msg_type == 'error':
+        msg_type = 'pyerr'
+    elif msg_type == 'execute_result':
+        msg_type = 'pyout'
+
     content = msg["content"]
     out = NotebookNode(output_type=msg_type)
 
     if msg_type == "stream":
         out.stream = content["name"]
-        out.text = content["data"]
+        try:
+            out.text = content['data']
+        except KeyError:
+            out.text = content['text']
     elif msg_type in ("display_data", "pyout"):
         for mime, data in content["data"].iteritems():
             if mime == "text/plain":
@@ -86,7 +111,7 @@ def run_cell(kernel_client, cell):
         if msg["msg_type"] == "status":
             if msg["content"]["execution_state"] == "idle" and input_acknowledged:
                 break
-        elif msg["msg_type"] == "pyin":
+        elif msg["msg_type"] in ('pyin', 'execute_input'):
             input_acknowledged = True
         else:
             out = parse_valid_output(msg)
@@ -111,7 +136,7 @@ def run_notebook(notebook):
             cell.outputs = run_cell(kernel_client, cell)
 
             cell.prompt_number = prompt_number
-            if cell.outputs:
+            if cell.outputs and cell.outputs[0]['output_type'] == 'pyout':
                 cell.outputs[0]["prompt_number"] = prompt_number
 
     kernel_manager.shutdown_kernel()