Merge pull request #6448 from takluyver/kernel-context-mgr

Add context manager for a kernel, and use it in nbconvert execute preprocessor

IPython/kernel/__init__.py

@@ -6,6 +6,6 b' from . import zmq'
 from .connect import *
 from .launcher import *
 from .client import KernelClient
-from .manager import KernelManager
+from .manager import KernelManager, run_kernel
 from .blocking import BlockingKernelClient
 from .multikernelmanager import MultiKernelManager

IPython/kernel/manager.py

@@ -5,12 +5,17 b''
 
 from __future__ import absolute_import
 
+from contextlib import contextmanager
 import os
 import re
 import signal
 import sys
 import time
 import warnings
+try:
+    from queue import Empty  # Py 3
+except ImportError:
+    from Queue import Empty  # Py 2
 
 import zmq
 
@@ -410,3 +415,39 b' class KernelManager(ConnectionFileMixin):'
 
 KernelManagerABC.register(KernelManager)
 
+
+def start_new_kernel(startup_timeout=60, kernel_name='python', **kwargs):
+    """Start a new kernel, and return its Manager and Client"""
+    km = KernelManager(kernel_name=kernel_name)
+    km.start_kernel(**kwargs)
+    kc = km.client()
+    kc.start_channels()
+
+    kc.kernel_info()
+    kc.get_shell_msg(block=True, timeout=startup_timeout)
+
+    # Flush channels
+    for channel in (kc.shell_channel, kc.iopub_channel):
+        while True:
+            try:
+                channel.get_msg(block=True, timeout=0.1)
+            except Empty:
+                break
+    return km, kc
+
+@contextmanager
+def run_kernel(**kwargs):
+    """Context manager to create a kernel in a subprocess.
+    
+    The kernel is shut down when the context exits.
+    
+    Returns
+    -------
+    kernel_client: connected KernelClient instance
+    """
+    km, kc = start_new_kernel(**kwargs)
+    try:
+        yield kc
+    finally:
+        kc.stop_channels()
+        km.shutdown_kernel(now=True)

IPython/kernel/tests/utils.py

@@ -15,7 +15,7 b' except ImportError:'
 import nose
 import nose.tools as nt
 
-from IPython.kernel import KernelManager
+from IPython.kernel import manager
 
 #-------------------------------------------------------------------------------
 # Globals
@@ -30,22 +30,13 b' KC = None'
 #-------------------------------------------------------------------------------
 # code
 #-------------------------------------------------------------------------------
-
-
-def start_new_kernel(argv=None):
-    """start a new kernel, and return its Manager and Client"""
-    km = KernelManager()
-    kwargs = dict(stdout=nose.iptest_stdstreams_fileno(), stderr=STDOUT)
-    if argv:
-        kwargs['extra_arguments'] = argv
-    km.start_kernel(**kwargs)
-    kc = km.client()
-    kc.start_channels()
+def start_new_kernel(**kwargs):
+    """start a new kernel, and return its Manager and Client
     
-    msg_id = kc.kernel_info()
-    kc.get_shell_msg(block=True, timeout=STARTUP_TIMEOUT)
-    flush_channels(kc)
-    return km, kc
+    Integrates with our output capturing for tests.
+    """
+    kwargs.update(dict(stdout=nose.iptest_stdstreams_fileno(), stderr=STDOUT))
+    return manager.start_new_kernel(startup_timeout=STARTUP_TIMEOUT, **kwargs)
 
 def flush_channels(kc=None):
     """flush any messages waiting on the queue"""
@@ -121,7 +112,6 b' def stop_global_kernel():'
     KM.shutdown_kernel(now=True)
     KM = None
 
-@contextmanager
 def new_kernel(argv=None):
     """Context manager for a new kernel in a subprocess
     
@@ -131,13 +121,11 b' def new_kernel(argv=None):'
     -------
     kernel_client: connected KernelClient instance
     """
-    km, kc = start_new_kernel(argv)
-    try:
-        yield kc
-    finally:
-        kc.stop_channels()
-        km.shutdown_kernel(now=True)
-
+    kwargs = dict(stdout=nose.iptest_stdstreams_fileno(), stderr=STDOUT,
+                  startup_timeout=STARTUP_TIMEOUT)
+    if argv is not None:
+        kwargs['extra_arguments'] = argv
+    return manager.run_kernel(**kwargs)
 
 def assemble_output(iopub):
     """assemble stdout/err from an execution"""

IPython/nbconvert/preprocessors/execute.py

@@ -46,38 +46,15 b' class ExecutePreprocessor(Preprocessor):'
     }
     
     extra_arguments = List(Unicode)
-    
-    def _create_client(self):
-        from IPython.kernel import KernelManager
-        self.km = KernelManager()
-        self.km.write_connection_file()
-        self.kc = self.km.client()
-        self.kc.start_channels()
-        self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
-        self.iopub = self.kc.iopub_channel
-        self.shell = self.kc.shell_channel
-        self.shell.kernel_info()
-        try:
-            self.shell.get_msg(timeout=self.timeout)
-        except Empty:
-            self.log.error("Timeout waiting for kernel_info reply")
-            raise
-        # flush IOPub
-        while True:
-            try:
-                self.iopub.get_msg(block=True, timeout=0.25)
-            except Empty:
-                break
-
-    def _shutdown_client(self):
-        self.kc.stop_channels()
-        self.km.shutdown_kernel()
-        del self.km
 
     def preprocess(self, nb, resources):
-        self._create_client()
-        nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
-        self._shutdown_client()
+        from IPython.kernel import run_kernel
+        kernel_name = nb.metadata.get('kernelspec', {}).get('name', 'python')
+        with run_kernel(kernel_name=kernel_name,
+                        extra_arguments=self.extra_arguments,
+                        stderr=open(os.devnull, 'w')) as kc:
+            self.kc = kc
+            nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
         return nb, resources
 
     def preprocess_cell(self, cell, resources, cell_index):
@@ -87,7 +64,7 b' class ExecutePreprocessor(Preprocessor):'
         if cell.cell_type != 'code':
             return cell, resources
         try:
-            outputs = self.run_cell(self.shell, self.iopub, cell)
+            outputs = self.run_cell(self.kc.shell_channel, self.kc.iopub_channel, cell)
         except Exception as e:
             self.log.error("failed to run cell: " + repr(e))
             self.log.error(str(cell.input))