Skip to content

Commit

Permalink
More type annotations
Browse files Browse the repository at this point in the history
  • Loading branch information
ianthomas23 committed Oct 2, 2024
1 parent d17c77f commit 5715dd0
Show file tree
Hide file tree
Showing 5 changed files with 21 additions and 17 deletions.
10 changes: 5 additions & 5 deletions ipykernel/kernelbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from datetime import datetime
from signal import SIGINT, SIGTERM, Signals

from .control import CONTROL_THREAD_NAME
from .thread import CONTROL_THREAD_NAME

if sys.platform != "win32":
from signal import SIGKILL
Expand Down Expand Up @@ -898,7 +898,7 @@ def kernel_info(self):
"supported_features": [],
}
if self._supports_kernel_subshells:
info["supported_features"].append("kernel subshells")
info["supported_features"] = ["kernel subshells"]
return info

async def kernel_info_request(self, socket, ident, parent):
Expand Down Expand Up @@ -1074,7 +1074,7 @@ async def do_debug_request(self, msg):
# Subshell control message handlers
# ---------------------------------------------------------------------------

async def create_subshell_request(self, socket, ident, parent):
async def create_subshell_request(self, socket, ident, parent) -> None:
if not self.session:
return
if not self._supports_kernel_subshells:
Expand All @@ -1089,7 +1089,7 @@ async def create_subshell_request(self, socket, ident, parent):

self.session.send(socket, "create_subshell_reply", reply, parent, ident)

async def delete_subshell_request(self, socket, ident, parent):
async def delete_subshell_request(self, socket, ident, parent) -> None:
if not self.session:
return
if not self._supports_kernel_subshells:
Expand All @@ -1111,7 +1111,7 @@ async def delete_subshell_request(self, socket, ident, parent):

self.session.send(socket, "delete_subshell_reply", reply, parent, ident)

async def list_subshell_request(self, socket, ident, parent):
async def list_subshell_request(self, socket, ident, parent) -> None:
if not self.session:
return
if not self._supports_kernel_subshells:
Expand Down
8 changes: 5 additions & 3 deletions ipykernel/shellchannel.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
"""A thread for a shell channel."""
import zmq.asyncio

from .subshell_manager import SubshellManager
from .thread import SHELL_CHANNEL_THREAD_NAME, BaseThread

Expand All @@ -9,21 +11,21 @@ class ShellChannelThread(BaseThread):
Communicates with shell/subshell threads via pairs of ZMQ inproc sockets.
"""

def __init__(self, context, shell_socket, **kwargs):
def __init__(self, context: zmq.asyncio.Context, shell_socket: zmq.asyncio.Socket, **kwargs):
"""Initialize the thread."""
super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs)
self._manager: SubshellManager | None = None
self._context = context
self._shell_socket = shell_socket

@property
def manager(self):
def manager(self) -> SubshellManager:
# Lazy initialisation.
if self._manager is None:
self._manager = SubshellManager(self._context, self._shell_socket)
return self._manager

def run(self):
def run(self) -> None:
"""Run the thread."""
try:
super().run()
Expand Down
4 changes: 2 additions & 2 deletions ipykernel/subshell.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def __init__(self, subshell_id: str, **kwargs):
# Inproc PAIR socket, for communication with shell channel thread.
self._pair_socket: zmq.asyncio.Socket | None = None

async def create_pair_socket(self, context: zmq.asyncio.Context, address: str):
async def create_pair_socket(self, context: zmq.asyncio.Context, address: str) -> None:
"""Create inproc PAIR socket, for communication with shell channel thread.
Should be called from this thread, so usually via add_task before the
Expand All @@ -27,7 +27,7 @@ async def create_pair_socket(self, context: zmq.asyncio.Context, address: str):
self._pair_socket = context.socket(zmq.PAIR)
self._pair_socket.connect(address)

def run(self):
def run(self) -> None:
try:
super().run()
finally:
Expand Down
10 changes: 6 additions & 4 deletions ipykernel/subshell_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ def list_subshell(self) -> list[str]:
with self._lock_cache:
return list(self._cache)

async def listen_from_control(self, subshell_task) -> None:
async def listen_from_control(self, subshell_task: t.Any) -> None:
"""Listen for messages on the control inproc socket, handle those messages and
return replies on the same socket. Runs in the shell channel thread.
"""
Expand All @@ -141,7 +141,7 @@ async def listen_from_subshells(self) -> None:
async for subshell_id in self._receive_stream:
tg.start_soon(self._listen_for_subshell_reply, subshell_id)

def subshell_id_from_thread_id(self, thread_id) -> str | None:
def subshell_id_from_thread_id(self, thread_id: int) -> str | None:
"""Return subshell_id of the specified thread_id.
Raises RuntimeError if thread_id is not the main shell or a subshell.
Expand Down Expand Up @@ -169,7 +169,7 @@ def _create_inproc_pair_socket(
socket.connect(address)
return socket

async def _create_subshell(self, subshell_task) -> str:
async def _create_subshell(self, subshell_task: t.Any) -> str:
"""Create and start a new subshell thread."""
assert current_thread().name == SHELL_CHANNEL_THREAD_NAME

Expand Down Expand Up @@ -241,7 +241,9 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None:
return
raise

async def _process_control_request(self, request, subshell_task) -> dict[str, t.Any]:
async def _process_control_request(
self, request: dict[str, t.Any], subshell_task: t.Any
) -> dict[str, t.Any]:
"""Process a control request message received on the control inproc
socket and return the reply. Runs in the shell channel thread.
"""
Expand Down
6 changes: 3 additions & 3 deletions ipykernel/thread.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@ def __init__(self, **kwargs):
self.pydev_do_not_trace = True
self.is_pydev_daemon_thread = True
self.__stop = Event()
self._tasks_and_args: t.List[t.Tuple[t.Callable, t.Tuple]] = []
self._tasks_and_args: t.List[t.Tuple[t.Any, t.Any]] = []

def add_task(self, task: t.Callable, *args: t.Tuple):
def add_task(self, task: t.Any, *args: t.Any) -> None:
# May only add tasks before the thread is started.
self._tasks_and_args.append((task, args))

Expand All @@ -34,7 +34,7 @@ async def _main(self) -> None:
await to_thread.run_sync(self.__stop.wait)
tg.cancel_scope.cancel()

def stop(self):
def stop(self) -> None:
"""Stop the thread.
This method is threadsafe.
Expand Down

0 comments on commit 5715dd0

Please sign in to comment.