text: import atexit
  import collections
  import datetime
  import base64
  import signal
  from textwrap import dedent
  
  try:
      from contextlib import asynccontextmanager
  except ImportError:
      # Use the backport package async-generator for Python < 3.7.
      # This should be removed when nbclient drops support for Python 3.6
      from async_generator import asynccontextmanager  # type: ignore
  from contextlib import contextmanager
  
  from time import monotonic
  from queue import Empty
  import asyncio
  import typing as t
  
  from traitlets.config.configurable import LoggingConfigurable
  from traitlets import List, Unicode, Bool, Enum, Any, Type, Dict, Integer, default
  
  from nbformat import NotebookNode
  from nbformat.v4 import output_from_msg
  from jupyter_client import KernelManager
  from jupyter_client.client import KernelClient
  
  from .exceptions import (
      CellControlSignal,
      CellTimeoutError,
      DeadKernelError,
      CellExecutionComplete,
      CellExecutionError
  )
  from .util import run_sync, ensure_async
  from .output_widget import OutputWidget
  
  
  def timestamp() -> str:
      return datetime.datetime.utcnow().isoformat() + 'Z'
  
  
  class NotebookClient(LoggingConfigurable):
      """
      Encompasses a Client for executing cells in a notebook
      """
  
      timeout: int = Integer(
          None,
          allow_none=True,
          help=dedent(
              """
              The time to wait (in seconds) for output from executions.
              If a cell execution takes longer, a TimeoutError is raised.
  
              ``None`` or ``-1`` will disable the timeout. If ``timeout_func`` is set,
              it overrides ``timeout``.
              """
          ),
      ).tag(config=True)
  
      timeout_func: t.Any = Any(
          default_value=None,
          allow_none=True,
          help=dedent(
              """
              A callable which, when given the cell source as input,
              returns the time to wait (in seconds) for output from cell
              executions. If a cell execution takes longer, a TimeoutError
              is raised.
  
              Returning ``None`` or ``-1`` will disable the timeout for the cell.
              Not setting ``timeout_func`` will cause the client to
              default to using the ``timeout`` trait for all cells. The
              ``timeout_func`` trait overrides ``timeout`` if it is not ``None``.
              """
          ),
      ).tag(config=True)
  
      interrupt_on_timeout: bool = Bool(
          False,
          help=dedent(
              """
              If execution of a cell times out, interrupt the kernel and
              continue executing other cells rather than throwing an error and
              stopping.
              """
          ),
      ).tag(config=True)
  
      startup_timeout: int = Integer(
          60,
          help=dedent(
              """
              The time to wait (in seconds) for the kernel to start.
              If kernel startup takes longer, a RuntimeError is
              raised.
              """
          ),
      ).tag(config=True)
  
      allow_errors: bool = Bool(
          False,
          help=dedent(
              """
              If ``False`` (default), when a cell raises an error the
              execution is stopped and a `CellExecutionError`
              is raised, except if the error name is in
              ``allow_error_names``.
              If ``True``, execution errors are ignored and the execution
              is continued until the end of the notebook. Output from
              exceptions is included in the cell output in both cases.
              """
          ),
      ).tag(config=True)
  
      allow_error_names: t.List[str] = List(
          Unicode(),
          help=dedent(
              """
              List of error names which won't stop the execution. Use this if the
              ``allow_errors`` option it too general and you want to allow only
              specific kinds of errors.
              """
          ),
      ).tag(config=True)
  
      force_raise_errors: bool = Bool(
          False,
          help=dedent(
              """
              If False (default), errors from executing the notebook can be
              allowed with a ``raises-exception`` tag on a single cell, or the
              ``allow_errors`` or ``allow_error_names`` configurable options for
              all cells. An allowed error will be recorded in notebook output, and
              execution will continue. If an error occurs when it is not
              explicitly allowed, a `CellExecutionError` will be raised.
              If True, `CellExecutionError` will be raised for any error that occurs
              while executing the notebook. This overrides the ``allow_errors``
              and ``allow_error_names`` options and the ``raises-exception`` cell
              tag.
              """
          ),
      ).tag(config=True)
  
      extra_arguments: t.List = List(Unicode()).tag(config=True)
  
      kernel_name: str = Unicode(
          '',
          help=dedent(
              """
              Name of kernel to use to execute the cells.
              If not set, use the kernel_spec embedded in the notebook.
              """
          ),
      ).tag(config=True)
  
      raise_on_iopub_timeout: bool = Bool(
          False,
          help=dedent(
              """
              If ``False`` (default), then the kernel will continue waiting for
              iopub messages until it receives a kernel idle message, or until a
              timeout occurs, at which point the currently executing cell will be
              skipped. If ``True``, then an error will be raised after the first
              timeout. This option generally does not need to be used, but may be
              useful in contexts where there is the possibility of executing
              notebooks with memory-consuming infinite loops.
              """
          ),
      ).tag(config=True)
  
      store_widget_state: bool = Bool(
          True,
          help=dedent(
              """
              If ``True`` (default), then the state of the Jupyter widgets created
              at the kernel will be stored in the metadata of the notebook.
              """
          ),
      ).tag(config=True)
  
      record_timing: bool = Bool(
          True,
          help=dedent(
              """
              If ``True`` (default), then the execution timings of each cell will
              be stored in the metadata of the notebook.
              """
          ),
      ).tag(config=True)
  
      iopub_timeout: int = Integer(
          4,
          allow_none=False,
          help=dedent(
              """
              The time to wait (in seconds) for IOPub output. This generally
              doesn't need to be set, but on some slow networks (such as CI
              systems) the default timeout might not be long enough to get all
              messages.
              """
          ),
      ).tag(config=True)
  
      shell_timeout_interval: int = Integer(
          5,
          allow_none=False,
          help=dedent(
              """
              The time to wait (in seconds) for Shell output before retrying.
              This generally doesn't need to be set, but if one needs to check
              for dead kernels at a faster rate this can help.
              """
          ),
      ).tag(config=True)
  
      shutdown_kernel = Enum(
          ['graceful', 'immediate'],
          default_value='graceful',
          help=dedent(
              """
              If ``graceful`` (default), then the kernel is given time to clean
              up after executing all cells, e.g., to execute its ``atexit`` hooks.
              If ``immediate``, then the kernel is signaled to immediately
              terminate.
              """
          ),
      ).tag(config=True)
  
      ipython_hist_file: str = Unicode(
          default_value=':memory:',
          help="""Path to file to use for SQLite history database for an IPython kernel.
  
          The specific value ``:memory:`` (including the colon
          at both end but not the back ticks), avoids creating a history file. Otherwise, IPython
          will create a history file for each kernel.
  
          When running kernels simultaneously (e.g. via multiprocessing) saving history a single
          SQLite file can result in database errors, so using ``:memory:`` is recommended in
          non-interactive contexts.
          """,
      ).tag(config=True)
  
      kernel_manager_class: KernelManager = Type(config=True, help='The kernel manager class to use.')
  
      @default('kernel_manager_class')
      def _kernel_manager_class_default(self) -> KernelManager:
          """Use a dynamic default to avoid importing jupyter_client at startup"""
          from jupyter_client import AsyncKernelManager
  
          return AsyncKernelManager
  
      _display_id_map: t.Dict[str, t.Dict] = Dict(
          help=dedent(
              """
                mapping of locations of outputs with a given display_id
                tracks cell index and output index within cell.outputs for
                each appearance of the display_id
                {
                     'display_id': {
                    cell_idx: [output_idx,]
                     }
                }
                """
          )
      )
  
      display_data_priority: t.List = List(
          [
              'text/html',
              'application/pdf',
              'text/latex',
              'image/svg+xml',
              'image/png',
              'image/jpeg',
              'text/markdown',
              'text/plain',
          ],
          help="""
              An ordered list of preferred output type, the first
              encountered will usually be used when converting discarding
              the others.
              """,
      ).tag(config=True)
  
      resources: t.Dict = Dict(
          help=dedent(
              """
              Additional resources used in the conversion process. For example,
              passing ``{'metadata': {'path': run_path}}`` sets the
              execution path to ``run_path``.
              """
          )
      )
  
      def __init__(
              self,
              nb: NotebookNode,
              km: t.Optional[KernelManager] = None,
              **kw) -> None:
          """Initializes the execution manager.
  
          Parameters
          ----------
          nb : NotebookNode
              Notebook being executed.
          km : KernelManager (optional)
              Optional kernel manager. If none is provided, a kernel manager will
              be created.
          """
          super().__init__(**kw)
          self.nb: NotebookNode = nb
          self.km: t.Optional[KernelManager] = km
          self.owns_km: bool = km is None  # whether the NotebookClient owns the kernel manager
          self.kc: t.Optional[KernelClient] = None
          self.reset_execution_trackers()
          self.widget_registry: t.Dict[str, t.Dict] = {
              '@jupyter-widgets/output': {
                  'OutputModel': OutputWidget
              }
          }
          # comm_open_handlers should return an object with a .handle_msg(msg) method or None
          self.comm_open_handlers: t.Dict[str, t.Any] = {
              'jupyter.widget': self.on_comm_open_jupyter_widget
          }
  
      def reset_execution_trackers(self) -> None:
          """Resets any per-execution trackers.
          """
          self.task_poll_for_reply: t.Optional[asyncio.Future] = None
          self.code_cells_executed = 0
          self._display_id_map = {}
          self.widget_state: t.Dict[str, t.Dict] = {}
          self.widget_buffers: t.Dict[str, t.List[t.Dict[str, str]]] = {}
          # maps to list of hooks, where the last is used, this is used
          # to support nested use of output widgets.
          self.output_hook_stack: t.Any = collections.defaultdict(list)
          # our front-end mimicing Output widgets
          self.comm_objects: t.Dict[str, t.Any] = {}
  
      def create_kernel_manager(self) -> KernelManager:
          """Creates a new kernel manager.
  
          Returns
          -------
          km : KernelManager
              Kernel manager whose client class is asynchronous.
          """
          if not self.kernel_name:
              kn = self.nb.metadata.get('kernelspec', {}).get('name')
              if kn is not None:
                  self.kernel_name = kn
  
          if not self.kernel_name:
              self.km = self.kernel_manager_class(config=self.config)
          else:
              self.km = self.kernel_manager_class(kernel_name=self.kernel_name, config=self.config)
  
          # If the current kernel manager is still using the default (synchronous) KernelClient class,
          # switch to the async version since that's what NBClient prefers.
          if self.km.client_class == 'jupyter_client.client.KernelClient':
              self.km.client_class = 'jupyter_client.asynchronous.AsyncKernelClient'
  
          return self.km
  
      async def _async_cleanup_kernel(self) -> None:
          assert self.km is not None
          now = self.shutdown_kernel == "immediate"
          try:
              # Queue the manager to kill the process, and recover gracefully if it's already dead.
              if await ensure_async(self.km.is_alive()):
                  await ensure_async(self.km.shutdown_kernel(now=now))
          except RuntimeError as e:
              # The error isn't specialized, so we have to check the message
              if 'No kernel is running!' not in str(e):
                  raise
          finally:
              # Remove any state left over even if we failed to stop the kernel
              await ensure_async(self.km.cleanup_resources())
              if getattr(self, "kc") and self.kc is not None:
                  await ensure_async(self.kc.stop_channels())
                  self.kc = None
                  self.km = None
  
      _cleanup_kernel = run_sync(_async_cleanup_kernel)
  
      async def async_start_new_kernel(self, **kwargs) -> None:
          """Creates a new kernel.
  
          Parameters
          ----------
          kwargs :
              Any options for ``self.kernel_manager_class.start_kernel()``. Because
              that defaults to AsyncKernelManager, this will likely include options
              accepted by ``AsyncKernelManager.start_kernel()``, which includes ``cwd``.
          """
          assert self.km is not None
          resource_path = self.resources.get('metadata', {}).get('path') or None
          if resource_path and 'cwd' not in kwargs:
              kwargs["cwd"] = resource_path
  
          has_history_manager_arg = any(
              arg.startswith('--HistoryManager.hist_file') for arg in self.extra_arguments)
          if (hasattr(self.km, 'ipykernel')
                  and self.km.ipykernel
                  and self.ipython_hist_file
                  and not has_history_manager_arg):
              self.extra_arguments += ['--HistoryManager.hist_file={}'.format(self.ipython_hist_file)]
  
          await ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs))
  
      start_new_kernel = run_sync(async_start_new_kernel)
  
      async def async_start_new_kernel_client(self) -> KernelClient:
          """Creates a new kernel client.
  
          Returns
          -------
          kc : KernelClient
              Kernel client as created by the kernel manager ``km``.
          """
          assert self.km is not None
          self.kc = self.km.client()
          await ensure_async(self.kc.start_channels())
          try:
              await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
          except RuntimeError:
              await self._async_cleanup_kernel()
              raise
          self.kc.allow_stdin = False
          return self.kc
  
      start_new_kernel_client = run_sync(async_start_new_kernel_client)
  
      @contextmanager
      def setup_kernel(self, **kwargs) -> t.Generator:
          """
          Context manager for setting up the kernel to execute a notebook.
  
          The assigns the Kernel Manager (``self.km``) if missing and Kernel Client(``self.kc``).
  
          When control returns from the yield it stops the client's zmq channels, and shuts
          down the kernel.
          """
          # by default, cleanup the kernel client if we own the kernel manager
          # and keep it alive if we don't
          cleanup_kc = kwargs.pop('cleanup_kc', self.owns_km)
  
          # Can't use run_until_complete on an asynccontextmanager function :(
          if self.km is None:
              self.km = self.create_kernel_manager()
  
          if not self.km.has_kernel:
              self.start_new_kernel(**kwargs)
              self.start_new_kernel_client()
          try:
              yield
          finally:
              if cleanup_kc:
                  self._cleanup_kernel()
  
      @asynccontextmanager
      async def async_setup_kernel(self, **kwargs) -> t.AsyncGenerator:
          """
          Context manager for setting up the kernel to execute a notebook.
  
          This assigns the Kernel Manager (``self.km``) if missing and Kernel Client(``self.kc``).
  
          When control returns from the yield it stops the client's zmq channels, and shuts
          down the kernel.
  
          Handlers for SIGINT and SIGTERM are also added to cleanup in case of unexpected shutdown.
          """
          # by default, cleanup the kernel client if we own the kernel manager
          # and keep it alive if we don't
          cleanup_kc = kwargs.pop('cleanup_kc', self.owns_km)
          if self.km is None:
              self.km = self.create_kernel_manager()
  
          # self._cleanup_kernel uses run_async, which ensures the ioloop is running again.
          # This is necessary as the ioloop has stopped once atexit fires.
          atexit.register(self._cleanup_kernel)
  
          def on_signal():
              asyncio.ensure_future(self._async_cleanup_kernel())
              atexit.unregister(self._cleanup_kernel)
  
          loop = asyncio.get_event_loop()
          try:
              loop.add_signal_handler(signal.SIGINT, on_signal)
              loop.add_signal_handler(signal.SIGTERM, on_signal)
          except (NotImplementedError, RuntimeError):
              # NotImplementedError: Windows does not support signals.
              # RuntimeError: Raised when add_signal_handler is called outside the main thread
              pass
  
          if not self.km.has_kernel:
              await self.async_start_new_kernel(**kwargs)
              await self.async_start_new_kernel_client()
          try:
              yield
          finally:
              if cleanup_kc:
                  await self._async_cleanup_kernel()
  
              atexit.unregister(self._cleanup_kernel)
              try:
                  loop.remove_signal_handler(signal.SIGINT)
                  loop.remove_signal_handler(signal.SIGTERM)
              except (NotImplementedError, RuntimeError):
                  pass
  
      async def async_execute(
              self,
              reset_kc: bool = False,
              **kwargs) -> NotebookNode:
          """
          Executes each code cell.
  
          Parameters
          ----------
          kwargs :
              Any option for ``self.kernel_manager_class.start_kernel()``. Because
              that defaults to AsyncKernelManager, this will likely include options
              accepted by ``jupyter_client.AsyncKernelManager.start_kernel()``,
              which includes ``cwd``.
  
              ``reset_kc`` if True, the kernel client will be reset and a new one
              will be created (default: False).
  
          Returns
          -------
          nb : NotebookNode
              The executed notebook.
          """
          if reset_kc and self.owns_km:
              await self._async_cleanup_kernel()
          self.reset_execution_trackers()
  
          async with self.async_setup_kernel(**kwargs):
              assert self.kc is not None
              self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
              msg_id = await ensure_async(self.kc.kernel_info())
              info_msg = await self.async_wait_for_reply(msg_id)
              if info_msg is not None:
                  if 'language_info' in info_msg['content']:
                      self.nb.metadata['language_info'] = info_msg['content']['language_info']
                  else:
                      raise RuntimeError(
                          'Kernel info received message content has no "language_info" key. '
                          'Content is:\n' + str(info_msg['content'])
                      )
              for index, cell in enumerate(self.nb.cells):
                  # Ignore `'execution_count' in content` as it's always 1
                  # when store_history is False
                  await self.async_execute_cell(
                      cell, index, execution_count=self.code_cells_executed + 1
                  )
              self.set_widgets_metadata()
  
          return self.nb
  
      execute = run_sync(async_execute)
  
      def set_widgets_metadata(self) -> None:
          if self.widget_state:
              self.nb.metadata.widgets = {
                  'application/vnd.jupyter.widget-state+json': {
                      'state': {
                          model_id: self._serialize_widget_state(state)
                          for model_id, state in self.widget_state.items()
                          if '_model_name' in state
                      },
                      'version_major': 2,
                      'version_minor': 0,
                  }
              }
              for key, widget in self.nb.metadata.widgets[
                  'application/vnd.jupyter.widget-state+json'
              ]['state'].items():
                  buffers = self.widget_buffers.get(key)
                  if buffers:
                      widget['buffers'] = buffers
  
      def _update_display_id(
              self,
              display_id: str,
              msg: t.Dict) -> None:
          """Update outputs with a given display_id"""
          if display_id not in self._display_id_map:
              self.log.debug("display id %r not in %s", display_id, self._display_id_map)
              return
  
          if msg['header']['msg_type'] == 'update_display_data':
              msg['header']['msg_type'] = 'display_data'
  
          try:
              out = output_from_msg(msg)
          except ValueError:
              self.log.error("unhandled iopub msg: " + msg['msg_type'])
              return
  
          for cell_idx, output_indices in self._display_id_map[display_id].items():
              cell = self.nb['cells'][cell_idx]
              outputs = cell['outputs']
              for output_idx in output_indices:
                  outputs[output_idx]['data'] = out['data']
                  outputs[output_idx]['metadata'] = out['metadata']
  
      async def _async_poll_for_reply(
              self,
              msg_id: str,
              cell: NotebookNode,
              timeout: t.Optional[int],
              task_poll_output_msg: asyncio.Future,
              task_poll_kernel_alive: asyncio.Future) -> t.Dict:
  
          assert self.kc is not None
          new_timeout: t.Optional[float] = None
          if timeout is not None:
              deadline = monotonic() + timeout
              new_timeout = float(timeout)
          while True:
              try:
                  msg = await ensure_async(self.kc.shell_channel.get_msg(timeout=new_timeout))
                  if msg['parent_header'].get('msg_id') == msg_id:
                      if self.record_timing:
                          cell['metadata']['execution']['shell.execute_reply'] = timestamp()
                      try:
                          await asyncio.wait_for(task_poll_output_msg, self.iopub_timeout)
                      except (asyncio.TimeoutError, Empty):
                          if self.raise_on_iopub_timeout:
                              task_poll_kernel_alive.cancel()
                              raise CellTimeoutError.error_from_timeout_and_cell(
                                  "Timeout waiting for IOPub output", self.iopub_timeout, cell
                              )
                          else:
                              self.log.warning("Timeout waiting for IOPub output")
                      task_poll_kernel_alive.cancel()
                      return msg
                  else:
                      if new_timeout is not None:
                          new_timeout = max(0, deadline - monotonic())
              except Empty:
                  # received no message, check if kernel is still alive
                  assert timeout is not None
                  task_poll_kernel_alive.cancel()
                  await self._async_check_alive()
                  await self._async_handle_timeout(timeout, cell)
  
      async def _async_poll_output_msg(
              self,
              parent_msg_id: str,
              cell: NotebookNode,
              cell_index: int) -> None:
  
          assert self.kc is not None
          while True:
              msg = await ensure_async(self.kc.iopub_channel.get_msg(timeout=None))
              if msg['parent_header'].get('msg_id') == parent_msg_id:
                  try:
                      # Will raise CellExecutionComplete when completed
                      self.process_message(msg, cell, cell_index)
                  except CellExecutionComplete:
                      return
  
      async def _async_poll_kernel_alive(self) -> None:
          while True:
              await asyncio.sleep(1)
              try:
                  await self._async_check_alive()
              except DeadKernelError:
                  assert self.task_poll_for_reply is not None
                  self.task_poll_for_reply.cancel()
                  return
  
      def _get_timeout(self, cell: t.Optional[NotebookNode]) -> int:
          if self.timeout_func is not None and cell is not None:
              timeout = self.timeout_func(cell)
          else:
              timeout = self.timeout
  
          if not timeout or timeout < 0:
              timeout = None
  
          return timeout
  
      async def _async_handle_timeout(
              self,
              timeout: int,
              cell: t.Optional[NotebookNode] = None) -> None:
  
          self.log.error("Timeout waiting for execute reply (%is)." % timeout)
          if self.interrupt_on_timeout:
              self.log.error("Interrupting kernel")
              assert self.km is not None
              await ensure_async(self.km.interrupt_kernel())
          else:
              raise CellTimeoutError.error_from_timeout_and_cell(
                  "Cell execution timed out", timeout, cell
              )
  
      async def _async_check_alive(self) -> None:
          assert self.kc is not None
          if not await ensure_async(self.kc.is_alive()):
              self.log.error("Kernel died while waiting for execute reply.")
              raise DeadKernelError("Kernel died")
  
      async def async_wait_for_reply(
              self,
              msg_id: str,
              cell: t.Optional[NotebookNode] = None) -> t.Optional[t.Dict]:
  
          assert self.kc is not None
          # wait for finish, with timeout
          timeout = self._get_timeout(cell)
          cummulative_time = 0
          while True:
              try:
                  msg = await ensure_async(
                      self.kc.shell_channel.get_msg(
                          timeout=self.shell_timeout_interval
                      )
                  )
              except Empty:
                  await self._async_check_alive()
                  cummulative_time += self.shell_timeout_interval
                  if timeout and cummulative_time > timeout:
                      await self._async_async_handle_timeout(timeout, cell)
                      break
              else:
                  if msg['parent_header'].get('msg_id') == msg_id:
                      return msg
          return None
  
      wait_for_reply = run_sync(async_wait_for_reply)
      # Backwards compatability naming for papermill
      _wait_for_reply = wait_for_reply
  
      def _passed_deadline(self, deadline: int) -> bool:
          if deadline is not None and deadline - monotonic() <= 0:
              return True
          return False
  
      def _check_raise_for_error(
              self,
              cell: NotebookNode,
              exec_reply: t.Optional[t.Dict]) -> None:
  
          if exec_reply is None:
              return None
  
          exec_reply_content = exec_reply['content']
          if exec_reply_content['status'] != 'error':
              return None
  
          cell_allows_errors = (not self.force_raise_errors) and (
              self.allow_errors
              or exec_reply_content.get('ename') in self.allow_error_names
              or "raises-exception" in cell.metadata.get("tags", []))
  
          if not cell_allows_errors:
              raise CellExecutionError.from_cell_and_msg(cell, exec_reply_content)
  
      async def async_execute_cell(
              self,
              cell: NotebookNode,
              cell_index: int,
              execution_count: t.Optional[int] = None,
              store_history: bool = True) -> NotebookNode:
          """
          Executes a single code cell.
  
          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.
          """
          assert self.kc is not None
          if cell.cell_type != 'code' or not cell.source.strip():
              self.log.debug("Skipping non-executing cell %s", cell_index)
              return cell
  
          if self.record_timing and 'execution' not in cell['metadata']:
              cell['metadata']['execution'] = {}
  
          self.log.debug("Executing cell:\n%s", cell.source)
  
          cell_allows_errors = (not self.force_raise_errors) and (
              self.allow_errors
              or "raises-exception" in cell.metadata.get("tags", []))
  
          parent_msg_id = await ensure_async(
              self.kc.execute(
                  cell.source,
                  store_history=store_history,
                  stop_on_error=not cell_allows_errors
              )
          )
          # We launched a code cell to execute
          self.code_cells_executed += 1
          exec_timeout = self._get_timeout(cell)
  
          cell.outputs = []
          self.clear_before_next_output = False
  
          task_poll_kernel_alive = asyncio.ensure_future(
              self._async_poll_kernel_alive()
          )
          task_poll_output_msg = asyncio.ensure_future(
              self._async_poll_output_msg(parent_msg_id, cell, cell_index)
          )
          self.task_poll_for_reply = asyncio.ensure_future(
              self._async_poll_for_reply(
                  parent_msg_id, cell, exec_timeout, task_poll_output_msg, task_poll_kernel_alive
              )
          )
          try:
              exec_reply = await self.task_poll_for_reply
          except asyncio.CancelledError:
              # can only be cancelled by task_poll_kernel_alive when the kernel is dead
              task_poll_output_msg.cancel()
              raise DeadKernelError("Kernel died")
          except Exception as e:
              # Best effort to cancel request if it hasn't been resolved
              try:
                  # Check if the task_poll_output is doing the raising for us
                  if not isinstance(e, CellControlSignal):
                      task_poll_output_msg.cancel()
              finally:
                  raise
  
          if execution_count:
              cell['execution_count'] = execution_count
          self._check_raise_for_error(cell, exec_reply)
          self.nb['cells'][cell_index] = cell
          return cell
  
      execute_cell = run_sync(async_execute_cell)
  
      def process_message(
              self,
              msg: t.Dict,
              cell: NotebookNode,
              cell_index: int) -> t.Optional[t.List]:
          """
          Processes a kernel message, updates cell state, and returns the
          resulting output object that was appended to cell.outputs.
  
          The input argument *cell* is modified in-place.
  
          Parameters
          ----------
          msg : dict
              The kernel message being processed.
          cell : nbformat.NotebookNode
              The cell which is currently being processed.
          cell_index : int
              The position of the cell within the notebook object.
  
          Returns
          -------
          output : dict
              The execution output payload (or None for no output).
  
          Raises
          ------
          CellExecutionComplete
            Once a message arrives which indicates computation completeness.
  
          """
          msg_type = msg['msg_type']
          self.log.debug("msg_type: %s", msg_type)
          content = msg['content']
          self.log.debug("content: %s", content)
  
          display_id = content.get('transient', {}).get('display_id', None)
          if display_id and msg_type in {'execute_result', 'display_data', 'update_display_data'}:
              self._update_display_id(display_id, msg)
  
          # set the prompt number for the input and the output
          if 'execution_count' in content:
              cell['execution_count'] = content['execution_count']
  
          if self.record_timing:
              if msg_type == 'status':
                  if content['execution_state'] == 'idle':
                      cell['metadata']['execution']['iopub.status.idle'] = timestamp()
                  elif content['execution_state'] == 'busy':
                      cell['metadata']['execution']['iopub.status.busy'] = timestamp()
              elif msg_type == 'execute_input':
                  cell['metadata']['execution']['iopub.execute_input'] = timestamp()
  
          if msg_type == 'status':
              if content['execution_state'] == 'idle':
                  raise CellExecutionComplete()
          elif msg_type == 'clear_output':
              self.clear_output(cell.outputs, msg, cell_index)
          elif msg_type.startswith('comm'):
              self.handle_comm_msg(cell.outputs, msg, cell_index)
          # Check for remaining messages we don't process
          elif msg_type not in ['execute_input', 'update_display_data']:
              # Assign output as our processed "result"
              return self.output(cell.outputs, msg, display_id, cell_index)
          return None
  
      def output(
              self,
              outs: t.List,
              msg: t.Dict,
              display_id: str,
              cell_index: int) -> t.Optional[t.List]:
  
          msg_type = msg['msg_type']
  
          parent_msg_id = msg['parent_header'].get('msg_id')
          if self.output_hook_stack[parent_msg_id]:
              # if we have a hook registered, it will overrride our
              # default output behaviour (e.g. OutputWidget)
              hook = self.output_hook_stack[parent_msg_id][-1]
              hook.output(outs, msg, display_id, cell_index)
              return None
  
          try:
              out = output_from_msg(msg)
          except ValueError:
              self.log.error("unhandled iopub msg: " + msg_type)
              return None
  
          if self.clear_before_next_output:
              self.log.debug('Executing delayed clear_output')
              outs[:] = []
              self.clear_display_id_mapping(cell_index)
              self.clear_before_next_output = False
  
          if display_id:
              # record output index in:
              #   _display_id_map[display_id][cell_idx]
              cell_map = self._display_id_map.setdefault(display_id, {})
              output_idx_list = cell_map.setdefault(cell_index, [])
              output_idx_list.append(len(outs))
  
          outs.append(out)
  
          return out
  
      def clear_output(
              self,
              outs: t.List,
              msg: t.Dict,
              cell_index: int) -> None:
  
          content = msg['content']
  
          parent_msg_id = msg['parent_header'].get('msg_id')
          if self.output_hook_stack[parent_msg_id]:
              # if we have a hook registered, it will overrride our
              # default clear_output behaviour (e.g. OutputWidget)
              hook = self.output_hook_stack[parent_msg_id][-1]
              hook.clear_output(outs, msg, cell_index)
              return
  
          if content.get('wait'):
              self.log.debug('Wait to clear output')
              self.clear_before_next_output = True
          else:
              self.log.debug('Immediate clear output')
              outs[:] = []
              self.clear_display_id_mapping(cell_index)
  
      def clear_display_id_mapping(
              self,
              cell_index: int) -> None:
  
          for display_id, cell_map in self._display_id_map.items():
              if cell_index in cell_map:
                  cell_map[cell_index] = []
  
      def handle_comm_msg(
              self,
              outs: t.List,
              msg: t.Dict,
              cell_index: int) -> None:
  
          content = msg['content']
          data = content['data']
          if self.store_widget_state and 'state' in data:  # ignore custom msg'es
              self.widget_state.setdefault(content['comm_id'], {}).update(data['state'])
              if 'buffer_paths' in data and data['buffer_paths']:
                  self.widget_buffers[content['comm_id']] = self._get_buffer_data(msg)
          # There are cases where we need to mimic a frontend, to get similar behaviour as
          # when using the Output widget from Jupyter lab/notebook
          if msg['msg_type'] == 'comm_open':
              target = msg['content'].get('target_name')
              handler = self.comm_open_handlers.get(target)
              if handler:
                  comm_id = msg['content']['comm_id']
                  comm_object = handler(msg)
                  if comm_object:
                      self.comm_objects[comm_id] = comm_object
              else:
                  self.log.warning(f'No handler found for comm target {target!r}')
          elif msg['msg_type'] == 'comm_msg':
              content = msg['content']
              comm_id = msg['content']['comm_id']
              if comm_id in self.comm_objects:
                  self.comm_objects[comm_id].handle_msg(msg)
  
      def _serialize_widget_state(self, state: t.Dict) -> t.Dict[str, t.Any]:
          """Serialize a widget state, following format in @jupyter-widgets/schema."""
          return {
              'model_name': state.get('_model_name'),
              'model_module': state.get('_model_module'),
              'model_module_version': state.get('_model_module_version'),
              'state': state,
          }
  
      def _get_buffer_data(self, msg: t.Dict) -> t.List[t.Dict[str, str]]:
          encoded_buffers = []
          paths = msg['content']['data']['buffer_paths']
          buffers = msg['buffers']
          for path, buffer in zip(paths, buffers):
              encoded_buffers.append(
                  {
                      'data': base64.b64encode(buffer).decode('utf-8'),
                      'encoding': 'base64',
                      'path': path,
                  }
              )
          return encoded_buffers
  
      def register_output_hook(
              self,
              msg_id: str,
              hook: OutputWidget) -> None:
          """Registers an override object that handles output/clear_output instead.
  
          Multiple hooks can be registered, where the last one will be used (stack based)
          """
          # mimics
          # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#registermessagehook
          self.output_hook_stack[msg_id].append(hook)
  
      def remove_output_hook(
              self,
              msg_id: str,
              hook: OutputWidget) -> None:
          """Unregisters an override object that handles output/clear_output instead"""
          # mimics
          # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#removemessagehook
          removed_hook = self.output_hook_stack[msg_id].pop()
          assert removed_hook == hook
  
      def on_comm_open_jupyter_widget(self, msg: t.Dict):
          content = msg['content']
          data = content['data']
          state = data['state']
          comm_id = msg['content']['comm_id']
          module = self.widget_registry.get(state['_model_module'])
          if module:
              widget_class = module.get(state['_model_name'])
              if widget_class:
                  return widget_class(comm_id, state, self.kc, self)
  
  
  def execute(
          nb: NotebookNode,
          cwd: t.Optional[str] = None,
          km: t.Optional[KernelManager] = None,
          **kwargs) -> NotebookClient:
      """Execute a notebook's code, updating outputs within the notebook object.
  
      This is a convenient wrapper around NotebookClient. It returns the
      modified notebook object.
  
      Parameters
      ----------
      nb : NotebookNode
        The notebook object to be executed
      cwd : str, optional
        If supplied, the kernel will run in this directory
      km : AsyncKernelManager, optional
        If supplied, the specified kernel manager will be used for code execution.
      kwargs :
        Any other options for NotebookClient, e.g. timeout, kernel_name
      """
      resources = {}
      if cwd is not None:
          resources['metadata'] = {'path': cwd}
      return NotebookClient(nb=nb, resources=resources, km=km, **kwargs).execute()
  