From 67c2f288b9950cbf2da4f58148baecb2e4db8775 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Thu, 14 May 2020 09:02:58 +0200
Subject: [PATCH] Support MultiKernelManager (#51)

---
 nbclient/client.py            | 24 ++++++++++++++++++----
 nbclient/tests/test_client.py | 38 +++++++++++++++++++++++++++++++++++
 requirements-dev.txt          |  1 +
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/nbclient/client.py b/nbclient/client.py
index 1bc46243..b021e133 100644
--- a/nbclient/client.py
+++ b/nbclient/client.py
@@ -363,17 +363,33 @@ async def async_start_new_kernel_client(self, **kwargs):
         -------
         kc : KernelClient
             Kernel client as created by the kernel manager `km`.
+        kernel_id : string-ized version 4 uuid
+           The id of the started kernel.
         """
         resource_path = self.resources.get('metadata', {}).get('path') or None
         if resource_path and 'cwd' not in kwargs:
             kwargs["cwd"] = resource_path
 
-        if self.km.ipykernel and self.ipython_hist_file:
+        # if self.km is a MultiKernelManager, it doesn't have an ipykernel attribute
+        # so no extra_arguments can be passed
+        if hasattr(self.km, 'ipykernel') and self.km.ipykernel and self.ipython_hist_file:
             self.extra_arguments += ['--HistoryManager.hist_file={}'.format(self.ipython_hist_file)]
 
-        await ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs))
+        kernel_id = await ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments,
+                                       **kwargs))
 
-        self.kc = self.km.client()
+        # if self.km is not a KernelManager, it's probably a MultiKernelManager
+        try:
+            self.km.client
+            km = self.km
+        except AttributeError:
+            try:
+                km = self.km.get_kernel(kernel_id)
+            except AttributeError:
+                raise AttributeError('self.km={} has no client() or get_kernel() method, '
+                                     'what is this?'.format(self.km))
+
+        self.kc = km.client()
         await ensure_async(self.kc.start_channels())
         try:
             await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
@@ -381,7 +397,7 @@ async def async_start_new_kernel_client(self, **kwargs):
             await self._async_cleanup_kernel()
             raise
         self.kc.allow_stdin = False
-        return self.kc
+        return self.kc, kernel_id
 
     start_new_kernel_client = run_sync(async_start_new_kernel_client)
 
diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py
index 9782f6cb..d95f864b 100644
--- a/nbclient/tests/test_client.py
+++ b/nbclient/tests/test_client.py
@@ -20,6 +20,7 @@
 import IPython
 from traitlets import TraitError
 from nbformat import NotebookNode
+from jupyter_client import KernelManager, MultiKernelManager
 from jupyter_client.kernelspec import KernelSpecManager
 from nbconvert.filters import strip_ansi
 from testpath import modified_env
@@ -400,6 +401,43 @@ def test_synchronous_setup_kernel():
     assert executor.kc is None
 
 
+def test_startnewkernel_with_kernelmanager():
+    nb = nbformat.v4.new_notebook()
+    km = KernelManager()
+    executor = NotebookClient(nb, km=km)
+    kc, kernel_id = executor.start_new_kernel_client()
+    # no kernel_id for a single kernel manager
+    assert kernel_id is None
+    # prove it initalized client
+    assert kc is not None
+    # since we are not using the setup_kernel context manager,
+    # cleanup has to be done manually
+    kc.shutdown()
+    km.cleanup()
+    kc.stop_channels()
+
+
+def test_startnewkernel_with_multikernelmanager():
+    nb = nbformat.v4.new_notebook()
+    km = MultiKernelManager()
+    executor = NotebookClient(nb, km=km)
+    kc, kernel_id = executor.start_new_kernel_client()
+    # a multi kernel manager always gives back an id to the started kernel
+    assert kernel_id is not None
+    # prove it initalized client
+    assert kc is not None
+    # since we are not using the setup_kernel context manager,
+    # cleanup has to be done manually
+    kc.shutdown()
+    km.cleanup(kernel_id)
+    km.remove_kernel(kernel_id)
+    # check that the kernel doesn't exist anymore
+    with pytest.raises(KeyError) as e_info:
+        km.get_kernel(kernel_id)
+    assert str(e_info.value) == "'Kernel with id not found: {}'".format(kernel_id)
+    kc.stop_channels()
+
+
 class TestExecute(NBClientTestsBase):
     """Contains test functions for execute.py"""
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 462b861c..c4c246b2 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,6 +1,7 @@
 codecov
 coverage
 ipython
+ipykernel
 ipywidgets
 pytest>=4.1
 pytest-cov>=2.6.1