Skip to content

Commit

Permalink
Update typing (#1150)
Browse files Browse the repository at this point in the history
  • Loading branch information
blink1073 authored Sep 30, 2023
1 parent 83ab7d0 commit 518785b
Show file tree
Hide file tree
Showing 11 changed files with 149 additions and 68 deletions.
1 change: 1 addition & 0 deletions ipykernel/comm/comm.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
if self.kernel is None:
self.kernel = Kernel.instance()

assert self.kernel.session is not None
self.kernel.session.send(
self.kernel.iopub_socket,
msg_type,
Expand Down
1 change: 1 addition & 0 deletions ipykernel/datapub.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def publish_data(self, data):
The data to be published. Think of it as a namespace.
"""
session = self.session
assert session is not None
buffers = serialize_object(
data,
buffer_threshold=session.buffer_threshold,
Expand Down
30 changes: 18 additions & 12 deletions ipykernel/displayhook.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@

# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.
from __future__ import annotations

import builtins
import sys
import typing as t

from IPython.core.displayhook import DisplayHook
from jupyter_client.session import Session, extract_header
Expand Down Expand Up @@ -61,35 +63,39 @@ class ZMQShellDisplayHook(DisplayHook):
session = Instance(Session, allow_none=True)
pub_socket = Any(allow_none=True)
parent_header = Dict({})
msg: dict[str, t.Any] | None

def set_parent(self, parent):
"""Set the parent for outbound messages."""
self.parent_header = extract_header(parent)

def start_displayhook(self):
"""Start the display hook."""
self.msg = self.session.msg(
"execute_result",
{
"data": {},
"metadata": {},
},
parent=self.parent_header,
)
if self.session:
self.msg = self.session.msg(
"execute_result",
{
"data": {},
"metadata": {},
},
parent=self.parent_header,
)

def write_output_prompt(self):
"""Write the output prompt."""
self.msg["content"]["execution_count"] = self.prompt_count
if self.msg:
self.msg["content"]["execution_count"] = self.prompt_count

def write_format_data(self, format_dict, md_dict=None):
"""Write format data to the message."""
self.msg["content"]["data"] = json_clean(encode_images(format_dict))
self.msg["content"]["metadata"] = md_dict
if self.msg:
self.msg["content"]["data"] = json_clean(encode_images(format_dict))
self.msg["content"]["metadata"] = md_dict

def finish_displayhook(self):
"""Finish up all displayhook activities."""
sys.stdout.flush()
sys.stderr.flush()
if self.msg["content"]["data"]:
if self.msg and self.msg["content"]["data"] and self.session:
self.session.send(self.pub_socket, self.msg, ident=self.topic)
self.msg = None
15 changes: 8 additions & 7 deletions ipykernel/inprocess/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,42 +60,43 @@ def _default_blocking_class(self):
def get_connection_info(self):
"""Get the connection info for the client."""
d = super().get_connection_info()
d["kernel"] = self.kernel
d["kernel"] = self.kernel # type:ignore[assignment]
return d

def start_channels(self, *args, **kwargs):
"""Start the channels on the client."""
super().start_channels()
self.kernel.frontends.append(self)
if self.kernel:
self.kernel.frontends.append(self)

@property
def shell_channel(self):
if self._shell_channel is None:
self._shell_channel = self.shell_channel_class(self)
self._shell_channel = self.shell_channel_class(self) # type:ignore[operator]
return self._shell_channel

@property
def iopub_channel(self):
if self._iopub_channel is None:
self._iopub_channel = self.iopub_channel_class(self)
self._iopub_channel = self.iopub_channel_class(self) # type:ignore[operator]
return self._iopub_channel

@property
def stdin_channel(self):
if self._stdin_channel is None:
self._stdin_channel = self.stdin_channel_class(self)
self._stdin_channel = self.stdin_channel_class(self) # type:ignore[operator]
return self._stdin_channel

@property
def control_channel(self):
if self._control_channel is None:
self._control_channel = self.control_channel_class(self)
self._control_channel = self.control_channel_class(self) # type:ignore[operator]
return self._control_channel

@property
def hb_channel(self):
if self._hb_channel is None:
self._hb_channel = self.hb_channel_class(self)
self._hb_channel = self.hb_channel_class(self) # type:ignore[operator]
return self._hb_channel

# Methods for sending specific messages
Expand Down
10 changes: 7 additions & 3 deletions ipykernel/inprocess/ipkernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ class InProcessKernel(IPythonKernel):
_underlying_iopub_socket = Instance(DummySocket, ())
iopub_thread: IOPubThread = Instance(IOPubThread) # type:ignore[assignment]

shell_stream = Instance(DummySocket, ())
shell_stream = Instance(DummySocket, ()) # type:ignore[arg-type]

@default("iopub_thread")
def _default_iopub_thread(self):
Expand All @@ -72,7 +72,8 @@ def __init__(self, **traits):
super().__init__(**traits)

self._underlying_iopub_socket.observe(self._io_dispatch, names=["message_sent"])
self.shell.kernel = self
if self.shell:
self.shell.kernel = self

async def execute_request(self, stream, ident, parent):
"""Override for temporary IO redirection."""
Expand All @@ -81,7 +82,8 @@ async def execute_request(self, stream, ident, parent):

def start(self):
"""Override registration of dispatchers for streams."""
self.shell.exit_now = False
if self.shell:
self.shell.exit_now = False

def _abort_queues(self):
"""The in-process kernel doesn't abort requests."""
Expand All @@ -99,6 +101,7 @@ def _input_request(self, prompt, ident, parent, password=False):

# Send the input request.
content = json_clean(dict(prompt=prompt, password=password))
assert self.session is not None
msg = self.session.msg("input_request", content, parent)
for frontend in self.frontends:
if frontend.session.session == parent["header"]["session"]:
Expand Down Expand Up @@ -132,6 +135,7 @@ def _redirected_io(self):
def _io_dispatch(self, change):
"""Called when a message is sent to the IO socket."""
assert self.iopub_socket.io_thread is not None
assert self.session is not None
ident, msg = self.session.recv(self.iopub_socket.io_thread.socket, copy=False)
for frontend in self.frontends:
frontend.iopub_channel.call_handlers(msg)
Expand Down
5 changes: 3 additions & 2 deletions ipykernel/inprocess/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,9 @@ def start_kernel(self, **kwds):

def shutdown_kernel(self):
"""Shutdown the kernel."""
self.kernel.iopub_thread.stop()
self._kill_kernel()
if self.kernel:
self.kernel.iopub_thread.stop()
self._kill_kernel()

def restart_kernel(self, now=False, **kwds):
"""Restart the kernel."""
Expand Down
29 changes: 21 additions & 8 deletions ipykernel/ipkernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ def __init__(self, **kwargs):
)

# Initialize the InteractiveShell subclass
self.shell = self.shell_class.instance(
self.shell = self.shell_class.instance( # type:ignore[attr-defined]
parent=self,
profile_dir=self.profile_dir,
user_module=self.user_module,
Expand Down Expand Up @@ -206,7 +206,8 @@ def dispatch_debugpy(self, msg):

@property
def banner(self):
return self.shell.banner
if self.shell:
return self.shell.banner

async def poll_stopped_queue(self):
"""Poll the stopped queue."""
Expand All @@ -215,7 +216,8 @@ async def poll_stopped_queue(self):

def start(self):
"""Start the kernel."""
self.shell.exit_now = False
if self.shell:
self.shell.exit_now = False
if self.debugpy_stream is None:
self.log.warning("debugpy_stream undefined, debugging will not be enabled")
else:
Expand All @@ -231,7 +233,7 @@ def set_parent(self, ident, parent, channel="shell"):
about the parent message.
"""
super().set_parent(ident, parent, channel)
if channel == "shell":
if channel == "shell" and self.shell:
self.shell.set_parent(parent)

def init_metadata(self, parent):
Expand Down Expand Up @@ -284,7 +286,8 @@ def _restore_input(self):

@property
def execution_count(self):
return self.shell.execution_count
if self.shell:
return self.shell.execution_count

@execution_count.setter
def execution_count(self, value):
Expand Down Expand Up @@ -348,6 +351,7 @@ async def do_execute(
):
"""Handle code execution."""
shell = self.shell # we'll need this a lot here
assert shell is not None

self._forward_input(allow_stdin)

Expand All @@ -371,7 +375,7 @@ async def run_cell(*args, **kwargs):
# not just asyncio
preprocessing_exc_tuple = None
try:
transformed_cell = self.shell.transform_cell(code)
transformed_cell = shell.transform_cell(code)
except Exception:
transformed_cell = code
preprocessing_exc_tuple = sys.exc_info()
Expand Down Expand Up @@ -488,6 +492,7 @@ def do_complete(self, code, cursor_pos):
cursor_pos = len(code)
line, offset = line_at_cursor(code, cursor_pos)
line_cursor = cursor_pos - offset
assert self.shell is not None
txt, matches = self.shell.complete("", line, line_cursor)
return {
"matches": matches,
Expand All @@ -509,6 +514,7 @@ def _experimental_do_complete(self, code, cursor_pos):
if cursor_pos is None:
cursor_pos = len(code)
with _provisionalcompleter():
assert self.shell is not None
raw_completions = self.shell.Completer.completions(code, cursor_pos)
completions = list(_rectify_completions(code, raw_completions))

Expand Down Expand Up @@ -548,6 +554,7 @@ def do_inspect(self, code, cursor_pos, detail_level=0, omit_sections=()):
reply_content: t.Dict[str, t.Any] = {"status": "ok"}
reply_content["data"] = {}
reply_content["metadata"] = {}
assert self.shell is not None
try:
if release.version_info >= (8,):
# `omit_sections` keyword will be available in IPython 8, see
Expand Down Expand Up @@ -581,6 +588,7 @@ def do_history(
unique=False,
):
"""Handle code history."""
assert self.shell is not None
if hist_access_type == "tail":
hist = self.shell.history_manager.get_tail(
n, raw=raw, output=output, include_latest=True
Expand All @@ -605,14 +613,16 @@ def do_history(

def do_shutdown(self, restart):
"""Handle kernel shutdown."""
self.shell.exit_now = True
if self.shell:
self.shell.exit_now = True
return dict(status="ok", restart=restart)

def do_is_complete(self, code):
"""Handle an is_complete request."""
transformer_manager = getattr(self.shell, "input_transformer_manager", None)
if transformer_manager is None:
# input_splitter attribute is deprecated
assert self.shell is not None
transformer_manager = self.shell.input_splitter
status, indent_spaces = transformer_manager.check_complete(code)
r = {"status": status}
Expand All @@ -628,6 +638,7 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):
from .serialize import serialize_object, unpack_apply_message

shell = self.shell
assert shell is not None
try:
working = shell.user_ns

Expand All @@ -652,6 +663,7 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):
for key in ns:
working.pop(key)

assert self.session is not None
result_buf = serialize_object(
result,
buffer_threshold=self.session.buffer_threshold,
Expand Down Expand Up @@ -686,7 +698,8 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):

def do_clear(self):
"""Clear the kernel."""
self.shell.reset(False)
if self.shell:
self.shell.reset(False)
return dict(status="ok")


Expand Down
2 changes: 1 addition & 1 deletion ipykernel/kernelapp.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,7 +547,7 @@ def init_kernel(self):
control_stream = ZMQStream(self.control_socket, self.control_thread.io_loop)
debugpy_stream = ZMQStream(self.debugpy_socket, self.control_thread.io_loop)
self.control_thread.start()
kernel_factory = self.kernel_class.instance
kernel_factory = self.kernel_class.instance # type:ignore[attr-defined]

kernel = kernel_factory(
parent=self,
Expand Down
Loading

0 comments on commit 518785b

Please sign in to comment.