Ensure KernelClient's API is complete

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]:

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)