Merge pull request #816 from mpacer/config_kernel_manager

Pass config to kernel from ExecutePreprocessor for configurable kernels

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,29 +211,71 @@ 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
         except ImportError:
             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