From a86c5b6b5e4b044aaa75561ab7d8d88e438a9da8 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Tue, 31 May 2022 18:38:15 +0200
Subject: [PATCH] Ensure KernelClient's API is complete

---
 jupyter_client/channels.py | 9 +++++----
 jupyter_client/client.py   | 8 ++++++++
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/jupyter_client/channels.py b/jupyter_client/channels.py
index c340c085e..76c43768d 100644
--- a/jupyter_client/channels.py
+++ b/jupyter_client/channels.py
@@ -9,7 +9,8 @@
 from threading import Event
 from threading import Thread
 
-import zmq.asyncio
+from zmq.asyncio import Context
+from zmq.asyncio import Socket
 
 from .channelsabc import HBChannelABC
 from .session import Session
@@ -49,7 +50,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 +194,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 +208,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..796a707f5 100644
--- a/jupyter_client/client.py
+++ b/jupyter_client/client.py
@@ -90,6 +90,14 @@ class KernelClient(ConnectionFileMixin):
     :meth:`get_shell_msg` to fetch messages from the shell channel.
     """
 
+    get_shell_msg: t.Union[t.Callable, t.Awaitable]
+    get_iopub_msg: t.Union[t.Callable, t.Awaitable]
+    get_stdin_msg: t.Union[t.Callable, t.Awaitable]
+    get_control_msg: t.Union[t.Callable, t.Awaitable]
+    wait_for_ready: t.Union[t.Callable, t.Awaitable]
+    is_alive: t.Union[t.Callable, t.Awaitable]
+    execute_interactive: t.Union[t.Callable, t.Awaitable]
+
     # The PyZMQ Context to use for communication with the kernel.
     context = Instance(zmq.asyncio.Context)