Change run_blocking into run_sync

Author: davidbrochart

nbclient/client.py

@@ -16,6 +16,7 @@
 from nbformat.v4 import output_from_msg
 
 from .exceptions import CellTimeoutError, DeadKernelError, CellExecutionComplete, CellExecutionError
+from .util import run_sync
 
 
 def timestamp():
@@ -294,43 +295,6 @@ def __init__(self, nb, km=None, **kw):
         self.km = km
         self.reset_execution_trackers()
 
-    def run_blocking(self, coro):
-        """Runs a coroutine and blocks until it has executed.
-
-        An event loop is created if no one already exists. If an event loop is
-        already running, this event loop execution is nested into the already
-        running one if `nest_asyncio` is set to True.
-
-        Parameters
-        ----------
-        coro : coroutine
-            The coroutine to be executed.
-
-        Returns
-        -------
-        result :
-            Whatever the coroutine returns.
-        """
-        try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-        if self.nest_asyncio:
-            import nest_asyncio
-            nest_asyncio.apply(loop)
-        try:
-            result = loop.run_until_complete(coro)
-        except RuntimeError as e:
-            if str(e) == 'This event loop is already running':
-                raise RuntimeError(
-                    'You are trying to run nbclient in an environment where an '
-                    'event loop is already running. Please pass `nest_asyncio=True` in '
-                    '`NotebookClient.execute` and such methods.'
-                )
-            raise
-        return result
-
     def reset_execution_trackers(self):
         """Resets any per-execution trackers.
         """
@@ -417,9 +381,6 @@ async def setup_kernel(self, **kwargs):
             self.kc.stop_channels()
             self.kc = None
 
-    def execute(self, **kwargs):
-        return self.run_blocking(self.async_execute(**kwargs))
-
     async def async_execute(self, **kwargs):
         """
         Executes each code cell.
@@ -445,6 +406,8 @@ async def async_execute(self, **kwargs):
 
         return self.nb
 
+    execute = run_sync(async_execute)
+
     def set_widgets_metadata(self):
         if self.widget_state:
             self.nb.metadata.widgets = {
@@ -511,7 +474,7 @@ async def _poll_for_reply(self, msg_id, cell, timeout, task_poll_output_msg):
                         timeout = max(0, deadline - monotonic())
             except Empty:
                 # received no message, check if kernel is still alive
-                self._check_alive()
+                await self._check_alive()
                 self._handle_timeout(timeout, cell)
 
     async def _poll_output_msg(self, parent_msg_id, cell, cell_index):
@@ -545,8 +508,8 @@ def _handle_timeout(self, timeout, cell=None):
                 "Cell execution timed out", timeout, cell
             )
 
-    def _check_alive(self):
-        if not self.kc.is_alive():
+    async def _check_alive(self):
+        if not await self.kc.is_alive():
             self.log.error("Kernel died while waiting for execute reply.")
             raise DeadKernelError("Kernel died")
 
@@ -559,7 +522,7 @@ async def _wait_for_reply(self, msg_id, cell=None):
             try:
                 msg = await self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
             except Empty:
-                self._check_alive()
+                await self._check_alive()
                 cummulative_time += self.shell_timeout_interval
                 if timeout and cummulative_time > timeout:
                     self._handle_timeout(timeout, cell)
@@ -591,11 +554,6 @@ def _check_raise_for_error(self, cell, exec_reply):
             if (exec_reply is not None) and exec_reply['content']['status'] == 'error':
                 raise CellExecutionError.from_cell_and_msg(cell, exec_reply['content'])
 
-    def execute_cell(self, cell, cell_index, execution_count=None, store_history=True):
-        return self.run_blocking(
-            self.async_execute_cell(cell, cell_index, execution_count, store_history)
-        )
-
     async def async_execute_cell(self, cell, cell_index, execution_count=None, store_history=True):
         """
         Executes a single code cell.
@@ -661,6 +619,8 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
         self.nb['cells'][cell_index] = cell
         return cell
 
+    execute_cell = run_sync(async_execute_cell)
+
     def process_message(self, msg, cell, cell_index):
         """
         Processes a kernel message, updates cell state, and returns the

nbclient/tests/test_client.py

@@ -672,6 +672,9 @@ def test_busy_message(self, executor, cell_mock, message_mock):
     )
     def test_deadline_exec_reply(self, executor, cell_mock, message_mock):
         # exec_reply is never received, so we expect to hit the timeout.
+        async def is_alive():
+            return True
+        executor.kc.is_alive = is_alive
         executor.kc.shell_channel.get_msg = MagicMock(side_effect=Empty())
         executor.timeout = 1
 

nbclient/util.py

@@ -0,0 +1,47 @@
+"""General utility methods"""
+
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+import asyncio
+
+
+def run_sync(coro):
+    """Runs a coroutine and blocks until it has executed.
+
+    An event loop is created if no one already exists. If an event loop is
+    already running, this event loop execution is nested into the already
+    running one if `nest_asyncio` is set to True.
+
+    Parameters
+    ----------
+    coro : coroutine
+        The coroutine to be executed.
+
+    Returns
+    -------
+    result :
+        Whatever the coroutine returns.
+    """
+    def wrapped(self, *args, **kwargs):
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        if self.nest_asyncio:
+            import nest_asyncio
+            nest_asyncio.apply(loop)
+        try:
+            result = loop.run_until_complete(coro(self, *args, **kwargs))
+        except RuntimeError as e:
+            if str(e) == 'This event loop is already running':
+                raise RuntimeError(
+                    'You are trying to run nbclient in an environment where an '
+                    'event loop is already running. Please pass `nest_asyncio=True` in '
+                    '`NotebookClient.execute` and such methods.'
+                )
+            raise
+        return result
+    wrapped.__doc__ = coro.__doc__
+    return wrapped

requirements.txt

@@ -1,5 +1,5 @@
 traitlets>=4.2
-jupyter_client>=6.0.0
+jupyter_client @ git+https://github.com/jupyter/jupyter_client
 nbformat>=5.0
 async_generator
 nest_asyncio