Convert to async/await and apply touch ups

Author: kevin-bates

notebook/notebookapp.py

@@ -7,6 +7,7 @@
 from __future__ import absolute_import, print_function
 
 import notebook
+import asyncio
 import binascii
 import datetime
 import errno
@@ -583,7 +584,7 @@ class NotebookApp(JupyterApp):
     flags = flags
 
     classes = [
-        KernelManager, Session, MappingKernelManager, KernelSpecManager,
+        KernelManager, Session, MappingKernelManager, AsyncMappingKernelManager, KernelSpecManager,
         ContentsManager, FileContentsManager, NotebookNotary,
         GatewayKernelManager, GatewayKernelSpecManager, GatewaySessionManager, GatewayClient,
     ]
@@ -1392,10 +1393,10 @@ def init_configurables(self):
         #  Ensure the appropriate jupyter_client is in place.
         if isinstance(self.kernel_manager, AsyncMappingKernelManager):
             if not async_kernel_mgmt_available:
-                raise ValueError("You're using `AsyncMappingKernelManager` without an appropriate "
+                raise ValueError("You are using `AsyncMappingKernelManager` without an appropriate "
                                  "jupyter_client installed!  Upgrade jupyter_client or change kernel managers.")
             else:
-                self.log.info("Asynchronous kernel management has been configured via '{}'.".
+                self.log.info("Asynchronous kernel management has been configured to use '{}'.".
                               format(self.kernel_manager.__class__.__name__))
 
         self.contents_manager = self.contents_manager_class(
@@ -1800,7 +1801,11 @@ def cleanup_kernels(self):
         n_kernels = len(self.kernel_manager.list_kernel_ids())
         kernel_msg = trans.ngettext('Shutting down %d kernel', 'Shutting down %d kernels', n_kernels)
         self.log.info(kernel_msg % n_kernels)
-        self.kernel_manager.shutdown_all()
+        # If we're using async kernel management, we need to invoke the async method via the event loop.
+        if isinstance(self.kernel_manager, AsyncMappingKernelManager):
+            asyncio.get_event_loop().run_until_complete(self.kernel_manager.shutdown_all())
+        else:
+            self.kernel_manager.shutdown_all()
 
     def notebook_info(self, kernel_count=True):
         "Return the current working directory and the server url information"

notebook/services/kernels/handlers.py

@@ -75,7 +75,7 @@ class KernelActionHandler(APIHandler):
     def post(self, kernel_id, action):
         km = self.kernel_manager
         if action == 'interrupt':
-            km.interrupt_kernel(kernel_id)
+            yield maybe_future(km.interrupt_kernel(kernel_id))
             self.set_status(204)
         if action == 'restart':
 

notebook/services/kernels/kernelmanager.py

@@ -401,8 +401,7 @@ def _handle_kernel_died(self, kernel_id):
         self.log.warning("Kernel %s died, removing from map.", kernel_id)
         self.remove_kernel(kernel_id)
 
-    @gen.coroutine
-    def start_kernel(self, kernel_id=None, path=None, **kwargs):
+    async def start_kernel(self, kernel_id=None, path=None, **kwargs):
         """Start a kernel for a session and return its kernel_id.
 
         Parameters
@@ -421,7 +420,7 @@ def start_kernel(self, kernel_id=None, path=None, **kwargs):
         if kernel_id is None:
             if path is not None:
                 kwargs['cwd'] = self.cwd_for_path(path)
-            kernel_id = yield super(AsyncMappingKernelManager, self).start_kernel(**kwargs)
+            kernel_id = await super(AsyncMappingKernelManager, self).start_kernel(**kwargs)
 
             self._kernel_connections[kernel_id] = 0
             self.start_watching_activity(kernel_id)
@@ -443,11 +442,9 @@ def start_kernel(self, kernel_id=None, path=None, **kwargs):
             self._check_kernel_id(kernel_id)
             self.log.info("Using existing kernel: %s" % kernel_id)
 
-        # py2-compat
-        raise gen.Return(kernel_id)
+        return kernel_id
 
-    @gen.coroutine
-    def shutdown_kernel(self, kernel_id, now=False, restart=False):
+    async def shutdown_kernel(self, kernel_id, now=False, restart=False):
         """Shutdown a kernel by kernel_id"""
         self._check_kernel_id(kernel_id)
         kernel = self._kernels[kernel_id]
@@ -464,13 +461,12 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
             type=self._kernels[kernel_id].kernel_name
         ).dec()
 
-        yield super(AsyncMappingKernelManager, self).shutdown_kernel(kernel_id, now=now, restart=restart)
+        await super(AsyncMappingKernelManager, self).shutdown_kernel(kernel_id, now=now, restart=restart)
 
-    @gen.coroutine
-    def restart_kernel(self, kernel_id, now=False):
+    async def restart_kernel(self, kernel_id, now=False):
         """Restart a kernel by kernel_id"""
         self._check_kernel_id(kernel_id)
-        yield super(AsyncMappingKernelManager, self).restart_kernel(kernel_id, now=now)
+        await super(AsyncMappingKernelManager, self).restart_kernel(kernel_id, now=now)
         kernel = self.get_kernel(kernel_id)
         # return a Future that will resolve when the kernel has successfully restarted
         channel = kernel.connect_shell()
@@ -506,7 +502,7 @@ def on_restart_failed():
         channel.on_recv(on_reply)
         loop = IOLoop.current()
         timeout = loop.add_timeout(loop.time() + self.kernel_info_timeout, on_timeout)
-        raise gen.Return(future)
+        return future
 
     def kernel_model(self, kernel_id):
         """Return a JSON-safe dict representing a kernel

notebook/services/sessions/tests/test_sessions_api.py

@@ -9,13 +9,22 @@
 import shutil
 import time
 
-pjoin = os.path.join
+from unittest import SkipTest
 
 from notebook.utils import url_path_join
 from notebook.tests.launchnotebook import NotebookTestBase, assert_http_error
 from nbformat.v4 import new_notebook
 from nbformat import write
 
+try:
+    from jupyter_client import AsyncMultiKernelManager
+    async_testing_enabled = True
+except ImportError:
+    async_testing_enabled = False
+
+pjoin = os.path.join
+
+
 class SessionAPI(object):
     """Wrapper for notebook API calls."""
     def __init__(self, request):
@@ -77,6 +86,7 @@ def modify_kernel_id(self, id, kernel_id):
     def delete(self, id):
         return self._req('DELETE', id)
 
+
 class SessionAPITest(NotebookTestBase):
     """Test the sessions web service API"""
     def setUp(self):
@@ -254,3 +264,24 @@ def test_modify_kernel_id(self):
         kernel.pop('last_activity')
         [ k.pop('last_activity') for k in kernel_list ]
         self.assertEqual(kernel_list, [kernel])
+
+
+class AsyncSessionAPITest(SessionAPITest):
+    """Test the sessions web service API using the AsyncMappingKernelManager"""
+
+    @classmethod
+    def get_argv(cls):
+        argv = super(AsyncSessionAPITest, cls).get_argv()
+
+        # before we extend the argv with the class, ensure that appropriate jupyter_client is available.
+        # if not available, don't set kernel_manager_class, resulting in the repeat of sync-based tests.
+        if async_testing_enabled:
+            argv.extend(['--NotebookApp.kernel_manager_class='
+                        'notebook.services.kernels.kernelmanager.AsyncMappingKernelManager'])
+        return argv
+
+    def setUp(self):
+        if not async_testing_enabled:
+            raise SkipTest("AsyncSessionAPITest.{test_method} skipped due to down-level jupyter_client!".
+                           format(test_method=self._testMethodName))
+        super(AsyncSessionAPITest, self).setUp()