client.py

"""nbclient implementation."""
from __future__ import annotations

import asyncio
import atexit
import base64
import collections
import datetime
import re
import signal
import typing as t
from contextlib import asynccontextmanager, contextmanager
from queue import Empty
from textwrap import dedent
from time import monotonic

from jupyter_client.client import KernelClient
from jupyter_client.manager import KernelManager
from nbformat import NotebookNode
from nbformat.v4 import output_from_msg
from traitlets import Any, Bool, Callable, Dict, Enum, Integer, List, Type, Unicode, default
from traitlets.config.configurable import LoggingConfigurable

from .exceptions import (
    CellControlSignal,
    CellExecutionComplete,
    CellExecutionError,
    CellTimeoutError,
    DeadKernelError,
)
from .output_widget import OutputWidget
from .util import ensure_async, run_hook, run_sync

_RGX_CARRIAGERETURN = re.compile(r".*\r(?=[^\n])")
_RGX_BACKSPACE = re.compile(r"[^\n]\b")

# mypy: disable-error-code="no-untyped-call"


def timestamp(msg: dict[str, t.Any] | None = None) -> str:
    """Get the timestamp for a message."""
    if msg and "header" in msg:  # The test mocks don't provide a header, so tolerate that
        msg_header = msg["header"]
        if "date" in msg_header and isinstance(msg_header["date"], datetime.datetime):
            try:
                # reformat datetime into expected format
                formatted_time = datetime.datetime.strftime(
                    msg_header["date"], "%Y-%m-%dT%H:%M:%S.%fZ"
                )
                if (
                    formatted_time
                ):  # docs indicate strftime may return empty string, so let's catch that too
                    return formatted_time
            except Exception:  # noqa
                pass  # fallback to a local time

    return datetime.datetime.utcnow().isoformat() + "Z"


class NotebookClient(LoggingConfigurable):
    """
    Encompasses a Client for executing cells in a notebook
    """

    timeout = 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.Callable[..., int | None] | None = Any(  # type:ignore[assignment]
        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(
        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)

    error_on_timeout = Dict(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            If a cell execution was interrupted after a timeout, don't wait for
            the execute_reply from the kernel (e.g. KeyboardInterrupt error).
            Instead, return an execute_reply with the given error, which should
            be of the following form::

                {
                    'ename': str,  # Exception name, as a string
                    'evalue': str,  # Exception value, as a string
                    'traceback': list(str),  # traceback frames, as strings
                }
            """
        ),
    ).tag(config=True)

    startup_timeout = 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(
        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 = 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(
        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)

    skip_cells_with_tag = Unicode(
        "skip-execution",
        help=dedent(
            """
            Name of the cell tag to use to denote a cell that should be skipped.
            """
        ),
    ).tag(config=True)

    extra_arguments = List(Unicode()).tag(config=True)

    kernel_name = 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(
        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(
        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(
        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 = 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 = 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 = 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 = Type(
        config=True, klass=KernelManager, help="The kernel manager class to use."
    )

    on_notebook_start = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes after the kernel manager and kernel client are setup, and
            cells are about to execute.
            Called with kwargs ``notebook``.
            """
        ),
    ).tag(config=True)

    on_notebook_complete = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes after the kernel is cleaned up.
            Called with kwargs ``notebook``.
            """
        ),
    ).tag(config=True)

    on_notebook_error = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes when the notebook encounters an error.
            Called with kwargs ``notebook``.
            """
        ),
    ).tag(config=True)

    on_cell_start = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes before a cell is executed and before non-executing cells
            are skipped.
            Called with kwargs ``cell`` and ``cell_index``.
            """
        ),
    ).tag(config=True)

    on_cell_execute = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes just before a code cell is executed.
            Called with kwargs ``cell`` and ``cell_index``.
            """
        ),
    ).tag(config=True)

    on_cell_complete = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes after a cell execution is complete. It is
            called even when a cell results in a failure.
            Called with kwargs ``cell`` and ``cell_index``.
            """
        ),
    ).tag(config=True)

    on_cell_executed = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes just after a code cell is executed, whether
            or not it results in an error.
            Called with kwargs ``cell``, ``cell_index`` and ``execute_reply``.
            """
        ),
    ).tag(config=True)

    on_cell_error = Callable(
        default_value=None,
        allow_none=True,
        help=dedent(
            """
            A callable which executes when a cell execution results in an error.
            This is executed even if errors are suppressed with ``cell_allows_errors``.
            Called with kwargs ``cell`, ``cell_index`` and ``execute_reply``.
            """
        ),
    ).tag(config=True)

    @default("kernel_manager_class")
    def _kernel_manager_class_default(self) -> type[KernelManager]:
        """Use a dynamic default to avoid importing jupyter_client at startup"""
        from jupyter_client import AsyncKernelManager  # type:ignore[attr-defined]

        return AsyncKernelManager

    _display_id_map: dict[str, t.Any] = Dict(  # type:ignore[assignment]
        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 = 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: dict[str, t.Any] = Dict(  # type:ignore[assignment]
        help=dedent(
            """
            Additional resources used in the conversion process. For example,
            passing ``{'metadata': {'path': run_path}}`` sets the
            execution path to ``run_path``.
            """
        )
    )

    coalesce_streams = Bool(
        help=dedent(
            """
            Merge all stream outputs with shared names into single streams.
            """
        )
    )

    def __init__(self, nb: NotebookNode, km: KernelManager | None = None, **kw: t.Any) -> 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: KernelManager | None = km
        self.owns_km: bool = km is None  # whether the NotebookClient owns the kernel manager
        self.kc: KernelClient | None = None
        self.reset_execution_trackers()
        self.widget_registry: dict[str, dict[str, t.Any]] = {
            "@jupyter-widgets/output": {"OutputModel": OutputWidget}
        }
        # comm_open_handlers should return an object with a .handle_msg(msg) method or None
        self.comm_open_handlers: 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: asyncio.Future[t.Any] | None = None
        self.code_cells_executed = 0
        self._display_id_map = {}
        self.widget_state: dict[str, dict[str, t.Any]] = {}
        self.widget_buffers: dict[str, dict[tuple[str, ...], 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 mimicking Output widgets
        self.comm_objects: 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)
        assert self.km is not None
        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", None) and self.kc is not None:
                await ensure_async(self.kc.stop_channels())  # type:ignore[func-returns-value]
                self.kc = None
                self.km = None

    _cleanup_kernel = run_sync(_async_cleanup_kernel)

    async def async_start_new_kernel(self, **kwargs: t.Any) -> 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 += [f"--HistoryManager.hist_file={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
        try:
            self.kc = self.km.client()
            await ensure_async(self.kc.start_channels())  # type:ignore[func-returns-value]
            await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
        except Exception as e:
            self.log.error(
                "Error occurred while starting new kernel client for kernel {}: {}".format(
                    getattr(self.km, "kernel_id", None), str(e)
                )
            )
            await self._async_cleanup_kernel()
            raise
        self.kc.allow_stdin = False
        await run_hook(self.on_notebook_start, notebook=self.nb)
        return self.kc

    start_new_kernel_client = run_sync(async_start_new_kernel_client)

    @contextmanager
    def setup_kernel(self, **kwargs: t.Any) -> t.Generator[None, None, None]:
        """
        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)

        if self.kc is None:
            self.start_new_kernel_client()

        try:
            yield
        finally:
            if cleanup_kc:
                self._cleanup_kernel()

    @asynccontextmanager
    async def async_setup_kernel(self, **kwargs: t.Any) -> t.AsyncGenerator[None, None]:
        """
        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() -> None:
            """Handle signals."""
            self._async_cleanup_kernel_future = 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 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)

        if self.kc is None:
            await self.async_start_new_kernel_client()

        try:
            yield
        except RuntimeError as e:
            await run_hook(self.on_notebook_error, notebook=self.nb)
            raise e
        finally:
            if cleanup_kc:
                await self._async_cleanup_kernel()
            await run_hook(self.on_notebook_complete, notebook=self.nb)
            atexit.unregister(self._cleanup_kernel)
            try:
                loop.remove_signal_handler(signal.SIGINT)
                loop.remove_signal_handler(signal.SIGTERM)
            except RuntimeError:
                pass

    async def async_execute(self, reset_kc: bool = False, **kwargs: t.Any) -> 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:
        """Set with widget metadata."""
        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"] = list(buffers.values())

    def _update_display_id(self, display_id: str, msg: dict[str, t.Any]) -> 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(f"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: int | None,
        task_poll_output_msg: asyncio.Future[t.Any],
        task_poll_kernel_alive: asyncio.Future[t.Any],
    ) -> dict[str, t.Any]:
        msg: dict[str, t.Any]
        assert self.kc is not None
        new_timeout: float | None = None
        if timeout is not None:
            deadline = monotonic() + timeout
            new_timeout = float(timeout)
        error_on_timeout_execute_reply = None
        while True:
            try:
                if error_on_timeout_execute_reply:
                    msg = error_on_timeout_execute_reply  # type:ignore[unreachable]
                    msg["parent_header"] = {"msg_id": msg_id}
                else:
                    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(msg)
                    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
                            ) from None
                        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()
                error_on_timeout_execute_reply = 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: NotebookNode | None) -> int | None:
        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: NotebookNode | None = None
    ) -> None | dict[str, t.Any]:
        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())
            if self.error_on_timeout:
                execute_reply = {"content": {**self.error_on_timeout, "status": "error"}}
                return execute_reply
            return None
        else:
            assert cell is not None
            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()):  # type:ignore[attr-defined]
            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: NotebookNode | None = None
    ) -> dict[str, t.Any] | None:
        """Wait for a message reply."""
        assert self.kc is not None
        # wait for finish, with timeout
        timeout = self._get_timeout(cell)
        cummulative_time = 0
        while True:
            try:
                msg: dict[str, t.Any] = 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_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 compatibility naming for papermill
    _wait_for_reply = wait_for_reply

    def _passed_deadline(self, deadline: int | None) -> bool:
        if deadline is not None and deadline - monotonic() <= 0:
            return True
        return False

    async def _check_raise_for_error(
        self, cell: NotebookNode, cell_index: int, exec_reply: dict[str, t.Any] | None
    ) -> 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", [])
        )
        await run_hook(
            self.on_cell_error, cell=cell, cell_index=cell_index, execute_reply=exec_reply
        )
        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: int | None = 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

        await run_hook(self.on_cell_start, cell=cell, cell_index=cell_index)

        if cell.cell_type != "code" or not cell.source.strip():
            self.log.debug("Skipping non-executing cell %s", cell_index)
            return cell

        if self.skip_cells_with_tag in cell.metadata.get("tags", []):
            self.log.debug("Skipping tagged cell %s", cell_index)
            return cell

        if self.record_timing:  # clear execution metadata prior to execution
            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", [])
        )

        await run_hook(self.on_cell_execute, cell=cell, cell_index=cell_index)
        parent_msg_id = await ensure_async(
            self.kc.execute(
                cell.source, store_history=store_history, stop_on_error=not cell_allows_errors
            )
        )
        await run_hook(self.on_cell_complete, cell=cell, cell_index=cell_index)
        # 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") from None
        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
        await run_hook(
            self.on_cell_executed, cell=cell, cell_index=cell_index, execute_reply=exec_reply
        )

        if self.coalesce_streams and cell.outputs:
            new_outputs = []
            streams: dict[str, NotebookNode] = {}
            for output in cell.outputs:
                if output["output_type"] == "stream":
                    if output["name"] in streams:
                        streams[output["name"]]["text"] += output["text"]
                    else:
                        new_outputs.append(output)
                        streams[output["name"]] = output
                else:
                    new_outputs.append(output)

            # process \r and \b characters
            for output in streams.values():
                old = output["text"]
                while len(output["text"]) < len(old):
                    old = output["text"]
                    # Cancel out anything-but-newline followed by backspace
                    output["text"] = _RGX_BACKSPACE.sub("", output["text"])
                # Replace all carriage returns not followed by newline
                output["text"] = _RGX_CARRIAGERETURN.sub("", output["text"])

            # We also want to ensure stdout and stderr are always in the same consecutive order,
            # because they are asynchronous, so order isn't guaranteed.
            for i, output in enumerate(new_outputs):
                if output["output_type"] == "stream" and output["name"] == "stderr":
                    if (
                        len(new_outputs) >= i + 2
                        and new_outputs[i + 1]["output_type"] == "stream"
                        and new_outputs[i + 1]["name"] == "stdout"
                    ):
                        stdout = new_outputs.pop(i + 1)
                        new_outputs.insert(i, stdout)

            cell.outputs = new_outputs

        await self._check_raise_for_error(cell, cell_index, exec_reply)

        self.nb["cells"][cell_index] = cell
        return cell

    execute_cell = run_sync(async_execute_cell)

    def process_message(
        self, msg: dict[str, t.Any], cell: NotebookNode, cell_index: int
    ) -> NotebookNode | None:
        """
        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 : NotebookNode
            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(msg)
                elif content["execution_state"] == "busy":
                    cell["metadata"]["execution"]["iopub.status.busy"] = timestamp(msg)
            elif msg_type == "execute_input":
                cell["metadata"]["execution"]["iopub.execute_input"] = timestamp(msg)

        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: list[NotebookNode], msg: dict[str, t.Any], display_id: str, cell_index: int
    ) -> NotebookNode | None:
        """Handle output."""

        msg_type = msg["msg_type"]
        out: NotebookNode | None = None

        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 override 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(f"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))

        if out:
            outs.append(out)

        return out

    def clear_output(
        self, outs: list[NotebookNode], msg: dict[str, t.Any], cell_index: int
    ) -> None:
        """Clear output."""
        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 override 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:
        """Clear a display id mapping for a cell."""
        for _, cell_map in self._display_id_map.items():
            if cell_index in cell_map:
                cell_map[cell_index] = []

    def handle_comm_msg(
        self, outs: list[NotebookNode], msg: dict[str, t.Any], cell_index: int
    ) -> None:
        """Handle a comm message."""
        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"]:
                comm_id = content["comm_id"]
                if comm_id not in self.widget_buffers:
                    self.widget_buffers[comm_id] = {}
                # for each comm, the path uniquely identifies a buffer
                new_buffers: dict[tuple[str, ...], dict[str, str]] = {
                    tuple(k["path"]): k for k in self._get_buffer_data(msg)
                }
                self.widget_buffers[comm_id].update(new_buffers)
        # 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: dict[str, t.Any]) -> 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: dict[str, t.Any]) -> list[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: dict[str, t.Any]) -> t.Any | None:
        """Handle a jupyter widget comm open."""
        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)
        return None


def execute(
    nb: NotebookNode,
    cwd: str | None = None,
    km: KernelManager | None = None,
    **kwargs: t.Any,
) -> NotebookNode:
    """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()