feat: preserve Jupyter MIME output in Streamlit

2026-08-31 00:52:34 +07:00
parent 6a8546cb0f
commit 0754cb2508
16 changed files with 1315 additions and 75 deletions
@@ -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