diff --git a/jupyter_client/channels.py b/jupyter_client/channels.py index c340c085e..a02228aeb 100644 --- a/jupyter_client/channels.py +++ b/jupyter_client/channels.py @@ -9,7 +9,9 @@ from threading import Event from threading import Thread -import zmq.asyncio +import zmq +from zmq.asyncio import Context +from zmq.asyncio import Socket from .channelsabc import HBChannelABC from .session import Session @@ -49,7 +51,7 @@ class HBChannel(Thread): def __init__( self, - context: t.Optional[zmq.asyncio.Context] = None, + context: t.Optional[Context] = None, session: t.Optional[Session] = None, address: t.Union[t.Tuple[str, int], str] = "", ): @@ -193,7 +195,7 @@ def call_handlers(self, since_last_heartbeat: float) -> None: class ZMQSocketChannel(object): """A ZMQ socket in an async API""" - def __init__(self, socket: zmq.asyncio.Socket, session: Session, loop: t.Any = None) -> None: + def __init__(self, socket: Socket, session: Session, loop: t.Any = None) -> None: """Create a channel. Parameters @@ -207,7 +209,7 @@ def __init__(self, socket: zmq.asyncio.Socket, session: Session, loop: t.Any = N """ super().__init__() - self.socket: t.Optional[zmq.asyncio.Socket] = socket + self.socket: t.Optional[Socket] = socket self.session = session async def _recv(self, **kwargs: t.Any) -> t.Dict[str, t.Any]: diff --git a/jupyter_client/client.py b/jupyter_client/client.py index 335aa2640..6fe78a67b 100644 --- a/jupyter_client/client.py +++ b/jupyter_client/client.py @@ -90,6 +90,26 @@ class KernelClient(ConnectionFileMixin): :meth:`get_shell_msg` to fetch messages from the shell channel. """ + get_shell_msg: t.Callable[ + ..., t.Union[t.Coroutine[t.Any, t.Any, t.Dict[str, t.Any]], t.Dict[str, t.Any]] + ] + get_iopub_msg: t.Callable[ + ..., t.Union[t.Coroutine[t.Any, t.Any, t.Dict[str, t.Any]], t.Dict[str, t.Any]] + ] + get_stdin_msg: t.Callable[ + ..., t.Union[t.Coroutine[t.Any, t.Any, t.Dict[str, t.Any]], t.Dict[str, t.Any]] + ] + get_control_msg: t.Callable[ + ..., t.Union[t.Coroutine[t.Any, t.Any, t.Dict[str, t.Any]], t.Dict[str, t.Any]] + ] + wait_for_ready: t.Callable[ + ["KernelClient", t.Optional[float]], t.Union[t.Coroutine[t.Any, t.Any, None], None] + ] + is_alive: t.Callable[["KernelClient"], t.Union[t.Coroutine[t.Any, t.Any, bool], bool]] + execute_interactive: t.Callable[ + ..., t.Union[t.Coroutine[t.Any, t.Any, t.Dict[str, t.Any]], t.Dict[str, t.Any]] + ] + # The PyZMQ Context to use for communication with the kernel. context = Instance(zmq.asyncio.Context)