lda/ipython-demo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: preserve Jupyter MIME output in Streamlit

Browse Source
This commit is contained in:
lda
2026-08-31 00:52:34 +07:00 Verified
parent 6a8546cb0f
commit 0754cb2508
16 changed files with 1315 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ipython_shell/jupyter/__init__.py
+29 -2
View File
@@ -1,11 +1,38 @@
from .app import JupyterApp
from .messages import (
DisplayDataMessage,
ErrorMessage,
ExecuteReplyMessage,
ExecuteResultMessage,
InputRequestMessage,
JupyterMessage,
KernelMessage,
KnownJupyterMessage,
ParsedJupyterMessage,
StatusMessage,
StreamMessage,
UnknownJupyterMessage,
parse_jupyter_message,
)
from .shell import JupyterShell
from .transport import InputRequest, JupyterTransport, KernelMessage
from .transport import InputRequest, JupyterTransport
__all__ = [
"InputRequest",
"DisplayDataMessage",
"ErrorMessage",
"ExecuteReplyMessage",
"ExecuteResultMessage",
"InputRequestMessage",
"JupyterApp",
"JupyterMessage",
"KernelMessage",
"JupyterShell",
"JupyterTransport",
"KernelMessage",
"KnownJupyterMessage",
"ParsedJupyterMessage",
"StatusMessage",
"StreamMessage",
"UnknownJupyterMessage",
"parse_jupyter_message",
]
src/ipython_shell/jupyter/app.py
+3 -2
View File
@@ -6,8 +6,9 @@ from uuid import uuid4
from ..models import AppInfo, ShellInfo
from ..registry import ShellRegistry
from ..utils import generate_call_id
from .messages import ParsedJupyterMessage
from .shell import JupyterShell
from .transport import InputRequest, KernelMessage
from .transport import InputRequest
class JupyterApp:
@@ -74,7 +75,7 @@ class JupyterApp:
self,
code: str,
shell_name: str = "last",
) -> AsyncIterator[KernelMessage | InputRequest]:
) -> AsyncIterator[ParsedJupyterMessage | InputRequest]:
"""Execute code and yield its subprocess messages."""
_, shell = self._select_shell(shell_name)
call_id = self._new_call_id()
src/ipython_shell/jupyter/messages.py
+260
View File
@@ -0,0 +1,260 @@
from typing import Annotated, ClassVar, Final, Literal
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
type JupyterChannel = Literal["iopub", "shell", "stdin"]
type MessageStatus = Literal["ok", "error", "abort"]
STREAM: Final[Literal["stream"]] = "stream"
DISPLAY_DATA: Final[Literal["display_data"]] = "display_data"
UPDATE_DISPLAY_DATA: Final[Literal["update_display_data"]] = "update_display_data"
EXECUTE_RESULT: Final[Literal["execute_result"]] = "execute_result"
ERROR: Final[Literal["error"]] = "error"
STATUS: Final[Literal["status"]] = "status"
EXECUTE_INPUT: Final[Literal["execute_input"]] = "execute_input"
CLEAR_OUTPUT: Final[Literal["clear_output"]] = "clear_output"
EXECUTE_REPLY: Final[Literal["execute_reply"]] = "execute_reply"
INPUT_REQUEST: Final[Literal["input_request"]] = "input_request"
class MimeContent(BaseModel):
"""The shared payload for rich display messages."""
data: dict[str, object]
metadata: dict[str, object] = Field(default_factory=dict)
transient: dict[str, object] | None = None
class ExecuteResultContent(MimeContent):
"""Rich display data plus the kernel execution counter."""
execution_count: int
class StreamContent(BaseModel):
"""A stdout or stderr chunk emitted by the kernel."""
name: Literal["stdout", "stderr"]
text: str
class ErrorContent(BaseModel):
"""An exception reported on the IOPub channel."""
ename: str
evalue: str
traceback: list[str]
class ExecuteReplyContent(BaseModel):
"""The terminal shell-channel reply for an execution request."""
status: MessageStatus
execution_count: int | None = None
ename: str | None = None
evalue: str | None = None
traceback: list[str] | None = None
class StatusContent(BaseModel):
"""Kernel busy/idle state from an IOPub status message."""
execution_state: Literal["busy", "idle", "starting"]
class InputRequestContent(BaseModel):
"""The prompt and echo policy for a kernel input request."""
prompt: str
password: bool = False
class ExecuteInputContent(BaseModel):
"""The code and execution counter echoed by the kernel."""
code: str
execution_count: int
class ClearOutputContent(BaseModel):
"""Whether a clear-output request should wait for new output."""
wait: bool = False
class JupyterMessage(BaseModel):
"""Common Jupyter envelope, also capable of holding unknown messages.
The protocol's message vocabulary is extensible, so ``msg_type`` stays a
string on this raw-compatible model. Known messages are narrowed into the
typed subclasses below by :func:`parse_jupyter_message`.
"""
model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")
channel: JupyterChannel = "iopub"
msg_type: str
message_id: str = ""
parent_id: str | None = None
metadata: dict[str, object] = Field(default_factory=dict)
content: dict[str, object] = Field(default_factory=dict)
buffers: list[bytes] = Field(default_factory=list)
class _TypedMessage(BaseModel):
"""Common fields for typed messages without narrowing mutable fields."""
model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")
message_id: str = ""
parent_id: str | None = None
metadata: dict[str, object] = Field(default_factory=dict)
buffers: list[bytes] = Field(default_factory=list)
class StreamMessage(_TypedMessage):
msg_type: Literal["stream"] = STREAM
channel: Literal["iopub"] = "iopub"
content: StreamContent
class DisplayDataMessage(_TypedMessage):
msg_type: Literal["display_data"] = DISPLAY_DATA
channel: Literal["iopub"] = "iopub"
content: MimeContent
class UpdateDisplayDataMessage(_TypedMessage):
msg_type: Literal["update_display_data"] = UPDATE_DISPLAY_DATA
channel: Literal["iopub"] = "iopub"
content: MimeContent
class ExecuteResultMessage(_TypedMessage):
msg_type: Literal["execute_result"] = EXECUTE_RESULT
channel: Literal["iopub"] = "iopub"
content: ExecuteResultContent
class ErrorMessage(_TypedMessage):
msg_type: Literal["error"] = ERROR
channel: Literal["iopub"] = "iopub"
content: ErrorContent
class StatusMessage(_TypedMessage):
msg_type: Literal["status"] = STATUS
channel: Literal["iopub"] = "iopub"
content: StatusContent
class ExecuteInputMessage(_TypedMessage):
msg_type: Literal["execute_input"] = EXECUTE_INPUT
channel: Literal["iopub"] = "iopub"
content: ExecuteInputContent
class ClearOutputMessage(_TypedMessage):
msg_type: Literal["clear_output"] = CLEAR_OUTPUT
channel: Literal["iopub"] = "iopub"
content: ClearOutputContent
class ExecuteReplyMessage(_TypedMessage):
msg_type: Literal["execute_reply"] = EXECUTE_REPLY
channel: Literal["shell"] = "shell"
content: ExecuteReplyContent
class InputRequestMessage(_TypedMessage):
msg_type: Literal["input_request"] = INPUT_REQUEST
channel: Literal["stdin"] = "stdin"
content: InputRequestContent
class UnknownJupyterMessage(JupyterMessage):
"""An extension or future message preserved without lossy parsing."""
raw: dict[str, object] = Field(default_factory=dict)
type KnownJupyterMessage = Annotated[
StreamMessage
| DisplayDataMessage
| UpdateDisplayDataMessage
| ExecuteResultMessage
| ErrorMessage
| StatusMessage
| ExecuteInputMessage
| ClearOutputMessage
| ExecuteReplyMessage
| InputRequestMessage,
Field(discriminator="msg_type"),
]
type ParsedJupyterMessage = KnownJupyterMessage | UnknownJupyterMessage
_KNOWN_MESSAGE_TYPES: Final[frozenset[str]] = frozenset(
{
STREAM,
DISPLAY_DATA,
UPDATE_DISPLAY_DATA,
EXECUTE_RESULT,
ERROR,
STATUS,
EXECUTE_INPUT,
CLEAR_OUTPUT,
EXECUTE_REPLY,
INPUT_REQUEST,
}
)
_KNOWN_MESSAGE_ADAPTER = TypeAdapter(KnownJupyterMessage)
def parse_jupyter_message(
raw: dict[str, object],
*,
channel: JupyterChannel,
) -> ParsedJupyterMessage:
"""Parse a decoded Jupyter message without dropping extensions."""
parent_header = raw.get("parent_header")
parent_id = (
parent_header.get("msg_id")
if isinstance(parent_header, dict)
and isinstance(parent_header.get("msg_id"), str)
else None
)
raw_metadata = raw.get("metadata")
metadata = raw_metadata if isinstance(raw_metadata, dict) else {}
raw_content = raw.get("content")
content = raw_content if isinstance(raw_content, dict) else {}
raw_buffers = raw.get("buffers")
buffers = (
[
bytes(buffer)
for buffer in raw_buffers
if isinstance(buffer, (bytes, bytearray, memoryview))
]
if isinstance(raw_buffers, list)
else []
)
raw_msg_type = raw.get("msg_type")
raw_message_id = raw.get("msg_id")
values: dict[str, object] = {
"channel": channel,
"msg_type": raw_msg_type if isinstance(raw_msg_type, str) else "",
"message_id": raw_message_id if isinstance(raw_message_id, str) else "",
"parent_id": parent_id,
"metadata": metadata,
"content": content,
"buffers": buffers,
}
if values["msg_type"] not in _KNOWN_MESSAGE_TYPES:
return UnknownJupyterMessage.model_validate({**values, "raw": raw})
return _KNOWN_MESSAGE_ADAPTER.validate_python(values)
# Existing callers use this name for the transport-level message. Keep it as
# an alias while the more expressive JupyterMessage name becomes canonical.
KernelMessage = JupyterMessage
src/ipython_shell/jupyter/shell.py
+14 -11
View File
@@ -1,10 +1,10 @@
from collections.abc import AsyncIterator
from dataclasses import replace
from datetime import UTC, datetime
from uuid import uuid4
from ..models import ShellInfo, ShellStatus
from .transport import InputRequest, JupyterTransport, KernelMessage
from .messages import ExecuteReplyMessage, ParsedJupyterMessage
from .transport import InputRequest, JupyterTransport
class JupyterShell:
@@ -30,7 +30,7 @@ class JupyterShell:
code: str,
*,
call_id: str,
) -> AsyncIterator[KernelMessage | InputRequest]:
) -> AsyncIterator[ParsedJupyterMessage | InputRequest]:
"""Yield kernel messages while one cell executes."""
await self.start()
self._active_call_id = call_id
@@ -39,19 +39,22 @@ class JupyterShell:
kernel_call_id = await self.transport.execute(code)
async for message in self.transport.messages_for(kernel_call_id):
if isinstance(message, InputRequest):
yield replace(message, call_id=call_id)
yield InputRequest(
call_id=call_id,
prompt=message.prompt,
password=message.password,
)
else:
# Keep the kernel's execution ID private to the transport.
public_message = replace(message, parent_id=call_id)
if public_message.msg_type == "execute_reply":
self._last_execution_count = int(
public_message.content.get(
"execution_count", self._last_execution_count
public_message = message.model_copy(update={"parent_id": call_id})
if isinstance(public_message, ExecuteReplyMessage):
if public_message.content.execution_count is not None:
self._last_execution_count = (
public_message.content.execution_count
)
)
self.status = (
"error"
if public_message.content.get("status") == "error"
if public_message.content.status == "error"
else "ready"
)
self.last_used_at = datetime.now(UTC).isoformat()
src/ipython_shell/jupyter/transport.py
+81 -31
View File
@@ -1,19 +1,22 @@
import asyncio
import os
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass
from time import monotonic
from typing import Any
from jupyter_client.manager import AsyncKernelManager
@dataclass
class KernelMessage:
"""A decoded Jupyter message kept at the transport boundary."""
msg_type: str
parent_id: str | None
content: dict[str, object]
buffers: list[bytes]
from .messages import (
ExecuteReplyMessage,
InputRequestMessage,
JupyterChannel,
JupyterMessage,
KernelMessage,
ParsedJupyterMessage,
StatusMessage,
parse_jupyter_message,
)
@dataclass
@@ -28,14 +31,21 @@ class InputRequest:
class JupyterTransport:
"""Own one persistent ipykernel subprocess and its message channels."""
def __init__(self) -> None:
_IOPUB_SETTLE_SECONDS = 0.25
def __init__(
self,
*,
matplotlib_backend: str | None = "module://matplotlib_inline.backend_inline",
) -> None:
self.manager = AsyncKernelManager()
self.matplotlib_backend = matplotlib_backend
self.client = None
self._active_call: str | None = None
self._call_lock = asyncio.Lock()
self._waiting_for_input: str | None = None
self._message_queue: (
asyncio.Queue[tuple[str, dict[str, Any] | BaseException]] | None
asyncio.Queue[tuple[JupyterChannel, dict[str, Any] | BaseException]] | None
) = None
self._reader_tasks: set[asyncio.Task[None]] = set()
@@ -44,7 +54,15 @@ class JupyterTransport:
if self.client is not None:
return
await self.manager.start_kernel()
# The parent process may select a non-interactive backend (Streamlit
# sets ``MPLBACKEND=Agg``). That is useful for server-side rendering,
# but it prevents ipykernel from publishing rich display MIME. Give
# each Jupyter subprocess its own inline default without mutating the
# host process environment.
kernel_env = dict(os.environ)
if self.matplotlib_backend is not None:
kernel_env["MPLBACKEND"] = self.matplotlib_backend
await self.manager.start_kernel(env=kernel_env)
self.client = self.manager.client()
self.client.start_channels()
await self.client.wait_for_ready()
@@ -67,7 +85,7 @@ class JupyterTransport:
async def _read_channel(
self,
channel: str,
channel: JupyterChannel,
get_message: Callable[..., Awaitable[dict[str, Any]]],
) -> None:
"""Read one ZMQ channel continuously into the transport queue."""
@@ -106,7 +124,7 @@ class JupyterTransport:
async def messages_for(
self, call_id: str
) -> AsyncIterator[KernelMessage | InputRequest]:
) -> AsyncIterator[ParsedJupyterMessage | InputRequest]:
"""Yield decoded output, input, and completion messages for one call."""
if self.client is None:
raise RuntimeError("JupyterTransport has not been started")
@@ -115,11 +133,11 @@ class JupyterTransport:
if self._message_queue is None:
raise RuntimeError("JupyterTransport message readers are not running")
reply: KernelMessage | None = None
reply: ExecuteReplyMessage | None = None
idle = False
try:
while True:
_, queued_message = await self._message_queue.get()
channel, queued_message = await self._message_queue.get()
if isinstance(queued_message, BaseException):
raise queued_message
@@ -128,37 +146,69 @@ class JupyterTransport:
if parent_id != call_id:
continue
if message["msg_type"] == "input_request":
decoded = parse_jupyter_message(message, channel=channel)
if isinstance(decoded, InputRequestMessage):
self._waiting_for_input = call_id
yield InputRequest(
call_id=call_id,
prompt=str(message.get("content", {}).get("prompt", "")),
password=bool(
message.get("content", {}).get("password", False)
),
prompt=decoded.content.prompt,
password=decoded.content.password,
)
continue
decoded = KernelMessage(
msg_type=message["msg_type"],
parent_id=parent_id,
content=dict(message.get("content", {})),
buffers=[bytes(buffer) for buffer in message.get("buffers", [])],
)
if decoded.msg_type == "execute_reply":
if isinstance(decoded, ExecuteReplyMessage):
# The shell reply and IOPub messages use different
# channels. Buffer the reply so it remains terminal
# even when both channels arrive close together.
reply = decoded
else:
yield decoded
if (
decoded.msg_type == "status"
and decoded.content.get("execution_state") == "idle"
if isinstance(decoded, StatusMessage) and (
decoded.content.execution_state == "idle"
):
idle = True
if reply is not None and idle:
# Shell and IOPub are independent ZMQ channels. The
# kernel can publish a final rich display after the shell
# reply and idle status have already reached this queue.
# Wait briefly for that cross-channel tail before ending
# the call, otherwise inline Matplotlib output can vanish.
settle_deadline = monotonic() + self._IOPUB_SETTLE_SECONDS
while True:
remaining = settle_deadline - monotonic()
if remaining <= 0:
break
try:
late_channel, late_message = await asyncio.wait_for(
self._message_queue.get(),
timeout=remaining,
)
except TimeoutError:
break
if isinstance(late_message, BaseException):
raise late_message
late_parent_id = late_message.get("parent_header", {}).get(
"msg_id"
)
if late_parent_id != call_id:
continue
late_decoded = parse_jupyter_message(
late_message,
channel=late_channel,
)
if isinstance(late_decoded, InputRequestMessage):
self._waiting_for_input = call_id
yield InputRequest(
call_id=call_id,
prompt=late_decoded.content.prompt,
password=late_decoded.content.password,
)
elif isinstance(late_decoded, ExecuteReplyMessage):
reply = late_decoded
else:
yield late_decoded
yield reply
return
finally: