Make nest_asyncio optional

Author: davidbrochart

binder/run_nbclient.ipynb

@@ -46,7 +46,7 @@
     "nb = nbf.read('./empty_notebook.ipynb', nbf.NO_CONVERT)\n",
     "\n",
     "# Execute our in-memory notebook, which will now have outputs\n",
-    "nb = nbclient.execute(nb)"
+    "nb = nbclient.execute(nb, nest_asyncio=True)"
    ]
   },
   {

nbclient/client.py

@@ -6,8 +6,6 @@
 # contextlib, and we `await yield_()` instead of just `yield`
 from async_generator import asynccontextmanager, async_generator, yield_
 
-import nest_asyncio
-
 from time import monotonic
 from queue import Empty
 import asyncio
@@ -97,6 +95,21 @@ class NotebookClient(LoggingConfigurable):
         ),
     ).tag(config=True)
 
+    nest_asyncio = Bool(
+        False,
+        help=dedent(
+            """
+            If False (default), then blocking functions such as `execute`
+            assume that no event loop is already running. These functions
+            run their async counterparts (e.g. `async_execute`) in an event
+            loop with `asyncio.run_until_complete`, which will fail if an
+            event loop is already running. This can be the case if nbclient
+            is used e.g. in a Jupyter Notebook. In that case, `nest_asyncio`
+            should be set to True.
+            """
+        ),
+    ).tag(config=True)
+
     force_raise_errors = Bool(
         False,
         help=dedent(
@@ -367,21 +380,9 @@ async def setup_kernel(self, **kwargs):
             self.kc.stop_channels()
             self.kc = None
 
-    def execute(self, **kwargs):
-        """
-        Executes each code cell (blocking).
-
-        Returns
-        -------
-        nb : NotebookNode
-            The executed notebook.
-        """
-        loop = get_loop()
-        return loop.run_until_complete(self.async_execute(**kwargs))
-
     async def async_execute(self, **kwargs):
         """
-        Executes each code cell asynchronously.
+        Executes each code cell.
 
         Returns
         -------
@@ -550,48 +551,9 @@ 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):
-        """
-        Executes a single code cell (blocking).
-
-        To execute all cells see :meth:`execute`.
-
-        Parameters
-        ----------
-        cell : nbformat.NotebookNode
-            The cell which is currently being processed.
-        cell_index : int
-            The position of the cell within the notebook object.
-        execution_count : int
-            The execution count to be assigned to the cell (default: Use kernel response)
-        store_history : bool
-            Determines if history should be stored in the kernel (default: False).
-            Specific to ipython kernels, which can store command histories.
-
-        Returns
-        -------
-        output : dict
-            The execution output payload (or None for no output).
-
-        Raises
-        ------
-        CellExecutionError
-            If execution failed and should raise an exception, this will be raised
-            with defaults about the failure.
-
-        Returns
-        -------
-        cell : NotebookNode
-            The cell which was just processed.
-        """
-        loop = get_loop()
-        return loop.run_until_complete(
-            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 asynchronously.
+        Executes a single code cell.
 
         To execute all cells see :meth:`execute`.
 
@@ -788,6 +750,24 @@ def _get_buffer_data(self, msg):
         return encoded_buffers
 
 
+def make_blocking(async_method):
+    def blocking_method(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)
+        return loop.run_until_complete(async_method(self, *args, **kwargs))
+    return blocking_method
+
+
+NotebookClient.execute = make_blocking(NotebookClient.async_execute)
+NotebookClient.execute_cell = make_blocking(NotebookClient.async_execute_cell)
+
+
 def execute(nb, cwd=None, km=None, **kwargs):
     """Execute a notebook's code, updating outputs within the notebook object.
 
@@ -809,15 +789,3 @@ def execute(nb, cwd=None, km=None, **kwargs):
     if cwd is not None:
         resources['metadata'] = {'path': cwd}
     return NotebookClient(nb=nb, resources=resources, km=km, **kwargs).execute()
-
-
-def get_loop():
-    try:
-        loop = asyncio.get_event_loop()
-    except RuntimeError:
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-    return loop
-
-
-nest_asyncio.apply()