You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
"""Implements an async kernel client"""
|
|
# Copyright (c) Jupyter Development Team.
|
|
# Distributed under the terms of the Modified BSD License.
|
|
from __future__ import annotations
|
|
|
|
import typing as t
|
|
|
|
import zmq.asyncio
|
|
from traitlets import Instance, Type
|
|
|
|
from ..channels import AsyncZMQSocketChannel, HBChannel
|
|
from ..client import KernelClient, reqrep
|
|
|
|
|
|
def wrapped(meth: t.Callable, channel: str) -> t.Callable:
|
|
"""Wrap a method on a channel and handle replies."""
|
|
|
|
def _(self: AsyncKernelClient, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
|
reply = kwargs.pop("reply", False)
|
|
timeout = kwargs.pop("timeout", None)
|
|
msg_id = meth(self, *args, **kwargs)
|
|
if not reply:
|
|
return msg_id
|
|
return self._recv_reply(msg_id, timeout=timeout, channel=channel)
|
|
|
|
return _
|
|
|
|
|
|
class AsyncKernelClient(KernelClient):
|
|
"""A KernelClient with async APIs
|
|
|
|
``get_[channel]_msg()`` methods wait for and return messages on channels,
|
|
raising :exc:`queue.Empty` if no message arrives within ``timeout`` seconds.
|
|
"""
|
|
|
|
context = Instance(zmq.asyncio.Context) # type:ignore[arg-type]
|
|
|
|
def _context_default(self) -> zmq.asyncio.Context:
|
|
self._created_context = True
|
|
return zmq.asyncio.Context()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Channel proxy methods
|
|
# --------------------------------------------------------------------------
|
|
|
|
get_shell_msg = KernelClient._async_get_shell_msg
|
|
get_iopub_msg = KernelClient._async_get_iopub_msg
|
|
get_stdin_msg = KernelClient._async_get_stdin_msg
|
|
get_control_msg = KernelClient._async_get_control_msg
|
|
|
|
wait_for_ready = KernelClient._async_wait_for_ready
|
|
|
|
# The classes to use for the various channels
|
|
shell_channel_class = Type(AsyncZMQSocketChannel) # type:ignore[arg-type]
|
|
iopub_channel_class = Type(AsyncZMQSocketChannel) # type:ignore[arg-type]
|
|
stdin_channel_class = Type(AsyncZMQSocketChannel) # type:ignore[arg-type]
|
|
hb_channel_class = Type(HBChannel) # type:ignore[arg-type]
|
|
control_channel_class = Type(AsyncZMQSocketChannel) # type:ignore[arg-type]
|
|
|
|
_recv_reply = KernelClient._async_recv_reply
|
|
|
|
# replies come on the shell channel
|
|
execute = reqrep(wrapped, KernelClient.execute)
|
|
history = reqrep(wrapped, KernelClient.history)
|
|
complete = reqrep(wrapped, KernelClient.complete)
|
|
is_complete = reqrep(wrapped, KernelClient.is_complete)
|
|
inspect = reqrep(wrapped, KernelClient.inspect)
|
|
kernel_info = reqrep(wrapped, KernelClient.kernel_info)
|
|
comm_info = reqrep(wrapped, KernelClient.comm_info)
|
|
|
|
is_alive = KernelClient._async_is_alive
|
|
execute_interactive = KernelClient._async_execute_interactive
|
|
|
|
# replies come on the control channel
|
|
shutdown = reqrep(wrapped, KernelClient.shutdown, channel="control")
|