diff --git a/nbconvert/preprocessors/execute.py b/nbconvert/preprocessors/execute.py index 8768328c1..ba7c1d8ff 100644 --- a/nbconvert/preprocessors/execute.py +++ b/nbconvert/preprocessors/execute.py @@ -5,18 +5,19 @@ # Distributed under the terms of the Modified BSD License. from textwrap import dedent +from contextlib import contextmanager try: from queue import Empty # Py 3 except ImportError: from Queue import Empty # Py 2 -from traitlets import List, Unicode, Bool, Enum, Any, Type, Dict, default +from traitlets import List, Unicode, Bool, Enum, Any, Type, Dict, Integer, default from nbformat.v4 import output_from_msg + from .base import Preprocessor from ..utils.exceptions import ConversionException -from traitlets import Integer class CellExecutionError(ConversionException): @@ -156,6 +157,17 @@ class ExecutePreprocessor(Preprocessor): """ ) ).tag(config=True) + + @default('kernel_name') + def _kernel_name_default(self): + try: + return self.nb.metadata.get('kernelspec', {}).get('name', 'python') + except AttributeError: + raise AttributeError('You did not specify a kernel_name for ' + 'the ExecutePreprocessor and you have not set ' + 'self.nb to be able to use that to infer the ' + 'kernel_name.') + raise_on_iopub_timeout = Bool(False, help=dedent( @@ -199,7 +211,7 @@ class ExecutePreprocessor(Preprocessor): help='The kernel manager class to use.' ) @default('kernel_manager_class') - def _km_default(self): + def _kernel_manager_class_default(self): """Use a dynamic default to avoid importing jupyter_client at startup""" try: from jupyter_client import KernelManager @@ -207,21 +219,63 @@ def _km_default(self): raise ImportError("`nbconvert --execute` requires the jupyter_client package: `pip install jupyter_client`") return KernelManager - # mapping of locations of outputs with a given display_id - # tracks cell index and output index within cell.outputs for - # each appearance of the display_id - # { - # 'display_id': { - # cell_idx: [output_idx,] - # } - # } - _display_id_map = Dict() + _display_id_map = Dict( + help=dedent( + """ + mapping of locations of outputs with a given display_id + tracks cell index and output index within cell.outputs for + each appearance of the display_id + { + 'display_id': { + cell_idx: [output_idx,] + } + } + """)) + + def start_new_kernel(self, **kwargs): + """Creates a new kernel manager and kernel client. - def preprocess(self, nb, resources): + Parameters + ---------- + kwargs : + Any options for `self.kernel_manager_class.start_kernel()`. Because + that defaults to KernelManager, this will likely include options + accepted by `KernelManager.start_kernel()``, which includes `cwd`. + + Returns + ------- + km : KernelManager + A kernel manager as created by self.kernel_manager_class. + kc : KernelClient + Kernel client as created by the kernel manager `km`. """ - Preprocess notebook executing each code cell. + km = self.kernel_manager_class(kernel_name=self.kernel_name, + config=self.config) + km.start_kernel(extra_arguments=self.extra_arguments, **kwargs) + + kc = km.client() + kc.start_channels() + try: + kc.wait_for_ready(timeout=self.startup_timeout) + except RuntimeError: + kc.stop_channels() + km.shutdown_kernel() + raise + kc.allow_stdin = False + return km, kc + + @contextmanager + def setup_preprocessor(self, nb, resources): + """ + Context manager for setting up the class to execute a notebook. - The input argument `nb` is modified in-place. + The assigns `nb` to `self.nb` where it will be modified in-place. It also creates + and assigns the Kernel Manager (`self.km`) and Kernel Client(`self.kc`). + + It is intended to yield to a block that will execute codeself. + + When control returns from the yield it stops the client's zmq channels, shuts + down the kernel, and removes the now unused attributes. Parameters ---------- @@ -239,48 +293,48 @@ def preprocess(self, nb, resources): resources : dictionary Additional resources used in the conversion process. """ - path = resources.get('metadata', {}).get('path', '') - if path == '': - path = None - + path = resources.get('metadata', {}).get('path', '') or None + self.nb = nb # clear display_id map self._display_id_map = {} - # from jupyter_client.manager import start_new_kernel - - def start_new_kernel(startup_timeout=60, kernel_name='python', **kwargs): - km = self.kernel_manager_class(kernel_name=kernel_name) - km.start_kernel(**kwargs) - kc = km.client() - kc.start_channels() - try: - kc.wait_for_ready(timeout=startup_timeout) - except RuntimeError: - kc.stop_channels() - km.shutdown_kernel() - raise - - return km, kc - - kernel_name = nb.metadata.get('kernelspec', {}).get('name', 'python') - if self.kernel_name: - kernel_name = self.kernel_name - self.log.info("Executing notebook with kernel: %s" % kernel_name) - self.km, self.kc = start_new_kernel( - startup_timeout=self.startup_timeout, - kernel_name=kernel_name, - extra_arguments=self.extra_arguments, - cwd=path) - self.kc.allow_stdin = False - self.nb = nb - + self.km, self.kc = self.start_new_kernel(cwd=path) try: - nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources) + # Yielding unbound args for more easier understanding and downstream consumption + yield nb, self.km, self.kc finally: self.kc.stop_channels() self.km.shutdown_kernel(now=self.shutdown_kernel == 'immediate') + + for attr in ['nb', 'km', 'kc']: + delattr(self, attr) + + def preprocess(self, nb, resources): + """ + Preprocess notebook executing each code cell. + + The input argument `nb` is modified in-place. - delattr(self, 'nb') + Parameters + ---------- + nb : NotebookNode + Notebook being executed. + resources : dictionary + Additional resources used in the conversion process. For example, + passing ``{'metadata': {'path': run_path}}`` sets the + execution path to ``run_path``. + + Returns + ------- + nb : NotebookNode + The executed notebook. + resources : dictionary + Additional resources used in the conversion process. + """ + + with self.setup_preprocessor(nb, resources): + self.log.info("Executing notebook with kernel: %s" % self.kernel_name) + nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources) return nb, resources