workload-automation/wlauto/utils/ipython.py

#    Copyright 2015 ARM Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import subprocess
from distutils.version import StrictVersion

import_error_str = ''
try:
    import IPython
except ImportError as import_error:
    IPython = None
    # Importing IPython can fail for a variety of reasons, report the actual
    # failure unless it's just that the package is not present
    if import_error.message.startswith("No module named"):  # pylint: disable=E1101
        import_error_str = 'ipynb_exporter requires ipython package to be installed'
    else:
        import_error_str = import_error.message

# The current code generates notebooks version 3
NBFORMAT_VERSION = 3


if IPython:
    if StrictVersion(IPython.__version__) >= StrictVersion('3.0.0'):
        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 StrictVersion(IPython.__version__) >= StrictVersion('2.0.0'):
        import IPython.kernel
        import IPython.nbformat.v3

        def read_notebook(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)  # pylint: disable=E1101

        NotebookNode = IPython.nbformat.v3.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT = ['ipython', 'nbconvert', '--to=latex', '--post=PDF']
    else:
        # Unsupported IPython version
        import_error_str = 'Unsupported IPython version {}'.format(IPython.__version__)


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"]
        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":
                attr = "text"
            else:
                attr = mime.split("/")[-1]
            setattr(out, attr, data)
    elif msg_type == "pyerr":
        out.ename = content["ename"]
        out.evalue = content["evalue"]
        out.traceback = content["traceback"]
    else:
        raise ValueError("Unknown msg_type {}".format(msg_type))

    return out


def run_cell(kernel_client, cell):
    """Run a cell of a notebook in an ipython kernel and return its output"""
    kernel_client.execute(cell.input)

    input_acknowledged = False
    outs = []
    while True:
        msg = kernel_client.get_iopub_msg()

        if msg["msg_type"] == "status":
            if msg["content"]["execution_state"] == "idle" and input_acknowledged:
                break
        elif msg["msg_type"] in ('pyin', 'execute_input'):
            input_acknowledged = True
        else:
            out = parse_valid_output(msg)
            outs.append(out)

    return outs


def run_notebook(notebook):
    """Run the notebook"""

    kernel_manager = IPython.kernel.KernelManager()
    kernel_manager.start_kernel(stderr=open(os.devnull, 'w'))
    kernel_client = kernel_manager.client()
    kernel_client.start_channels()

    for sheet in notebook.worksheets:
        for (prompt_number, cell) in enumerate(sheet.cells, 1):
            if cell.cell_type != "code":
                continue

            cell.outputs = run_cell(kernel_client, cell)

            cell.prompt_number = prompt_number
            if cell.outputs and cell.outputs[0]['output_type'] == 'pyout':
                cell.outputs[0]["prompt_number"] = prompt_number

    kernel_manager.shutdown_kernel()


def generate_pdf(nbbasename, output_directory):
    """Generate a PDF from the ipython notebook

    ipython nbconvert claims that the CLI is not stable, so keep this
    function here to be able to cope with inconsistencies

    """

    prev_dir = os.getcwd()
    os.chdir(output_directory)

    with open(os.devnull, 'w') as devnull:
        subprocess.check_call(IPYTHON_NBCONVERT + [nbbasename], stderr=devnull)

    os.chdir(prev_dir)