From abd651ee7831414e64dfc6c4ff03d5e2d9d72e2b Mon Sep 17 00:00:00 2001 From: seem Date: Wed, 10 Apr 2024 17:44:04 +0200 Subject: [PATCH] render plots inline in notebook sessions Plus some minor general housekeeping. --- .../positron/positron_ipykernel/plots.py | 117 ++++++++++-------- .../positron_ipykernel/positron_comm.py | 22 +++- .../positron_ipykernel/positron_ipkernel.py | 33 ++--- .../positron_ipykernel/session_mode.py | 27 ++++ .../positron_ipykernel/tests/conftest.py | 5 +- .../positron_ipykernel/tests/test_plots.py | 64 ++++++---- .../tests/test_positron_ipkernel.py | 7 +- .../positron/positron_ipykernel/widget.py | 6 +- .../positron/positron_language_server.py | 5 +- .../browser/positronPlotsService.ts | 49 ++++---- 10 files changed, 205 insertions(+), 130 deletions(-) create mode 100644 extensions/positron-python/python_files/positron/positron_ipykernel/session_mode.py diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py b/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py index e8375e1fa8d..9ecae76f816 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py @@ -8,14 +8,15 @@ import logging import pickle import uuid -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import comm from .plot_comm import PlotBackendMessageContent, PlotResult, RenderRequest from .positron_comm import CommMessage, JsonRpcErrorCode, PositronComm +from .session_mode import SessionMode from .utils import JsonRecord -from .widget import _WIDGET_MIME_TYPE +from .widget import WIDGET_MIME_TYPE logger = logging.getLogger(__name__) @@ -27,42 +28,58 @@ class PositronDisplayPublisherHook: - def __init__(self, target_name: str): + def __init__(self, target_name: str, session_mode: SessionMode): + self.target_name = target_name + self.session_mode = session_mode + self.comms: Dict[str, PositronComm] = {} self.figures: Dict[str, str] = {} - self.target_name = target_name self.fignums: List[int] = [] - def __call__(self, msg, *args, **kwargs) -> Optional[dict]: - if msg["msg_type"] == "display_data": - # If there is no image for our display, don't create a - # positron.plot comm and let the parent deal with the msg. - data = msg["content"]["data"] - if _WIDGET_MIME_TYPE in data: - # This is a widget, let the widget hook handle it - return msg - if "image/png" not in data: - return msg - - # Otherwise, try to pickle the current figure so that we - # can restore the context for future renderings. We construct - # a new plot comm to advise the client of the new figure. - pickled = self._pickle_current_figure() - if pickled is not None: - id = str(uuid.uuid4()) - self.figures[id] = pickled - - # Creating a comm per plot figure allows the client - # to request new renderings of each plot at a later time, - # e.g. on resizing the plots view - self._create_comm(id) - - # Returning None implies our hook has processed the message - # and it stops the parent from sending the display_data via - # the standard iopub channel - return None + def __call__(self, msg: Dict[str, Any]) -> Optional[Dict[str, Any]]: + # The display publisher calls each hook on the message in the order they were registered. + # If a hook returns a message, that message is passed to the next hook, and eventually sent + # to the frontend. If a hook returns None, no further hooks are called and the message is not + # sent to the frontend. + + if self.session_mode == SessionMode.NOTEBOOK: + # We're in a notebook session, let the notebook UI handle the display + return msg + + if msg["msg_type"] != "display_data": + # It's not a display_data message, do nothing + return msg + + data = msg["content"]["data"] + + if WIDGET_MIME_TYPE in data: + # This is a widget, let the widget hook handle it + return msg + + if "image/png" not in data: + # There is no attached png image, do nothing + return msg + + # Otherwise, try to pickle the current figure so that we + # can restore the context for future renderings. We construct + # a new plot comm to advise the client of the new figure. + pickled = self._pickle_current_figure() + if pickled is None: + logger.warning("No figure ") + return msg + + id = str(uuid.uuid4()) + self.figures[id] = pickled - return msg + # Creating a comm per plot figure allows the client + # to request new renderings of each plot at a later time, + # e.g. on resizing the plots view + self._create_comm(id) + + # Returning None implies our hook has processed the message + # and it stops the parent from sending the display_data via + # the standard iopub channel + return None def _create_comm(self, comm_id: str) -> None: """ @@ -120,34 +137,32 @@ def shutdown(self) -> None: # -- Private Methods -- def _pickle_current_figure(self) -> Optional[str]: - pickled = None - figure = None - # Delay importing matplotlib until the kernel and shell has been initialized # otherwise the graphics backend will be reset to the gui import matplotlib.pyplot as plt # We turn off interactive mode before accessing the plot context - was_interactive = plt.isinteractive() - plt.ioff() - - # Check to see if there are any figures left in stack to display - # If not, get the number of figures to display from matplotlib - if len(self.fignums) == 0: - self.fignums = plt.get_fignums() + with plt.ioff(): + # Check to see if there are any figures left in stack to display + # If not, get the number of figures to display from matplotlib + if len(self.fignums) == 0: + self.fignums = plt.get_fignums() + + if len(self.fignums) == 0: + logger.warning("Hook called without a figure to display") + return None - # Get the current figure, remove from it from being called next hook - if len(self.fignums) > 0: + # Get the current figure, remove it from displayed in the next call figure = plt.figure(self.fignums.pop(0)) - # Pickle the current figure - if figure is not None and not self._is_figure_empty(figure): - pickled = codecs.encode(pickle.dumps(figure), "base64").decode() + if self._is_figure_empty(figure): + logger.warning("Figure is empty") + return None - if was_interactive: - plt.ion() + # Pickle the current figure + pickled = codecs.encode(pickle.dumps(figure), "base64").decode() - return pickled + return pickled def _resize_pickled_figure( self, diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/positron_comm.py b/extensions/positron-python/python_files/positron/positron_ipykernel/positron_comm.py index 8cedb73df5f..d534430f7bc 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/positron_comm.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/positron_comm.py @@ -10,7 +10,14 @@ import comm -from . import connections_comm, data_explorer_comm, help_comm, plot_comm, ui_comm, variables_comm +from . import ( + connections_comm, + data_explorer_comm, + help_comm, + plot_comm, + ui_comm, + variables_comm, +) from ._vendor.pydantic import ValidationError from ._vendor.pydantic.generics import GenericModel from .utils import JsonData, JsonRecord @@ -21,10 +28,23 @@ ## Create an enum of JSON-RPC error codes @enum.unique class JsonRpcErrorCode(enum.IntEnum): + # Documentation below is taken directly from https://www.jsonrpc.org/specification#error_object + # for convenience. + + # Invalid JSON was received by the server. + # An error occurred on the server while parsing the JSON text. PARSE_ERROR = -32700 + + # The JSON sent is not a valid Request object. INVALID_REQUEST = -32600 + + # The method does not exist / is not available. METHOD_NOT_FOUND = -32601 + + # Invalid method parameter(s). INVALID_PARAMS = -32602 + + # Internal JSON-RPC error. INTERNAL_ERROR = -32603 diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/positron_ipkernel.py b/extensions/positron-python/python_files/positron/positron_ipykernel/positron_ipkernel.py index 01b81885ead..c6b1d6e3653 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/positron_ipkernel.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/positron_ipkernel.py @@ -20,7 +20,13 @@ from ipykernel.zmqshell import ZMQDisplayPublisher, ZMQInteractiveShell from IPython.core import oinspect, page from IPython.core.interactiveshell import ExecutionInfo, InteractiveShell -from IPython.core.magic import Magics, MagicsManager, line_magic, magics_class, needs_local_scope +from IPython.core.magic import ( + Magics, + MagicsManager, + line_magic, + magics_class, + needs_local_scope, +) from IPython.utils import PyColorize from .connections import ConnectionsService @@ -28,6 +34,7 @@ from .help import HelpService, help from .lsp import LSPService from .plots import PositronDisplayPublisherHook +from .session_mode import SessionMode from .ui import UiService from .utils import JsonRecord from .variables import VariablesService @@ -45,26 +52,6 @@ class _CommTarget(str, enum.Enum): Connections = "positron.connection" -class SessionMode(str, enum.Enum): - """ - The mode that the kernel application was started in. - """ - - Console = "console" - Notebook = "notebook" - Background = "background" - - Default = Console - - def __str__(self) -> str: - # Override for better display in argparse help. - return self.value - - @classmethod - def trait(cls) -> traitlets.Enum: - return traitlets.Enum(sorted(cls), help=cls.__doc__) - - logger = logging.getLogger(__name__) @@ -282,7 +269,7 @@ def _showtraceback(self, etype, evalue: Exception, stb: List[str]): # type: ign """ Enhance tracebacks for the Positron frontend. """ - if self.session_mode == SessionMode.Notebook: + if self.session_mode == SessionMode.NOTEBOOK: # Don't modify the traceback in a notebook. The frontend assumes that it's unformatted # and applies its own formatting. return super()._showtraceback(etype, evalue, stb) # type: ignore IPython type annotation is wrong @@ -356,7 +343,7 @@ def __init__(self, **kwargs) -> None: # Create Positron services self.data_explorer_service = DataExplorerService(_CommTarget.DataExplorer) - self.display_pub_hook = PositronDisplayPublisherHook(_CommTarget.Plot) + self.display_pub_hook = PositronDisplayPublisherHook(_CommTarget.Plot, self.session_mode) self.ui_service = UiService() self.help_service = HelpService() self.lsp_service = LSPService(self) diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/session_mode.py b/extensions/positron-python/python_files/positron/positron_ipykernel/session_mode.py new file mode 100644 index 00000000000..48e1213d426 --- /dev/null +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/session_mode.py @@ -0,0 +1,27 @@ +# +# Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved. +# + +import enum + +import traitlets + + +class SessionMode(str, enum.Enum): + """ + The mode that the kernel application was started in. + """ + + CONSOLE = "console" + NOTEBOOK = "notebook" + BACKGROUND = "background" + + DEFAULT = CONSOLE + + def __str__(self) -> str: + # Override for better display in argparse help. + return self.value + + @classmethod + def trait(cls) -> traitlets.Enum: + return traitlets.Enum(sorted(cls), help=cls.__doc__) diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/conftest.py b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/conftest.py index 66bbb6affc9..e82cf1a506c 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/conftest.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/conftest.py @@ -8,14 +8,15 @@ import comm import pytest from traitlets.config import Config + from positron_ipykernel.connections import ConnectionsService from positron_ipykernel.data_explorer import DataExplorerService from positron_ipykernel.positron_ipkernel import ( PositronIPKernelApp, PositronIPyKernel, PositronShell, - SessionMode, ) +from positron_ipykernel.session_mode import SessionMode from positron_ipykernel.variables import VariablesService @@ -55,7 +56,7 @@ def kernel() -> PositronIPyKernel: app.config = Config() # Needed to avoid traitlets errors # Positron-specific attributes: - app.session_mode = SessionMode.Console + app.session_mode = SessionMode.CONSOLE kernel = PositronIPyKernel.instance(parent=app) diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_plots.py b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_plots.py index 5e100974a2d..8f978066790 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_plots.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_plots.py @@ -7,7 +7,7 @@ import io import pickle from pathlib import Path -from typing import Iterable, cast +from typing import Iterable, Optional, cast import matplotlib import matplotlib.pyplot as plt @@ -17,13 +17,17 @@ from matplotlib.figure import Figure from matplotlib.testing.compare import compare_images from matplotlib_inline.backend_inline import configure_inline_support + from positron_ipykernel.plots import BASE_DPI, PositronDisplayPublisherHook from positron_ipykernel.positron_comm import JsonRpcErrorCode +from positron_ipykernel.session_mode import SessionMode +from positron_ipykernel.utils import JsonRecord from .conftest import DummyComm, PositronShell -from .utils import comm_request, json_rpc_error, json_rpc_request +from .utils import comm_open_message, comm_request, json_rpc_error, json_rpc_request PLOT_DATA = [1, 2] +TARGET_NAME = "target_name" @pytest.fixture(autouse=True) @@ -55,24 +59,46 @@ def images_path() -> Path: @pytest.fixture def hook() -> PositronDisplayPublisherHook: - return PositronDisplayPublisherHook("positron.plot") + return PositronDisplayPublisherHook(TARGET_NAME, SessionMode.CONSOLE) @pytest.fixture -def figure_comm(hook: PositronDisplayPublisherHook) -> DummyComm: +def notebook_hook() -> PositronDisplayPublisherHook: + return PositronDisplayPublisherHook(TARGET_NAME, SessionMode.NOTEBOOK) + + +def display_data_message() -> JsonRecord: """ - A comm corresponding to a test figure belonging to the Positron display publisher hook. + A valid display_data message with an image/png MIME type. """ + # The display hook doesn't depend on the image/png value, so ignore it. + return comm_request({"image/png": None}, msg_type="display_data") + + +def init_hook(hook: PositronDisplayPublisherHook) -> Optional[JsonRecord]: # Initialize the hook by calling it on a figure created with the test plot data plt.plot(PLOT_DATA) - msg = comm_request({"image/png": None}, msg_type="display_data") - hook(msg) - plt.close() + try: + msg = display_data_message() + return hook(msg) + finally: + plt.close() + + +@pytest.fixture +def figure_comm(hook: PositronDisplayPublisherHook) -> DummyComm: + """ + A comm corresponding to a test figure belonging to the Positron display publisher hook. + """ + assert init_hook(hook) is None # Return the comm corresponding to the first figure id = next(iter(hook.comms)) figure_comm = cast(DummyComm, hook.comms[id].comm) + # Check that the comm_open message was sent + assert figure_comm.messages == [comm_open_message(TARGET_NAME)] + # Clear messages due to the comm_open figure_comm.messages.clear() @@ -93,11 +119,15 @@ def test_hook_call_noop_on_no_image_png(hook: PositronDisplayPublisherHook) -> N assert hook.comms == {} +def test_hook_call_noop_in_notebook(notebook_hook: PositronDisplayPublisherHook) -> None: + msg = display_data_message() + assert notebook_hook(msg) == msg + assert notebook_hook.figures == {} + assert notebook_hook.comms == {} + + def test_hook_call(hook: PositronDisplayPublisherHook, images_path: Path) -> None: - # It returns `None` to indicate that it's consumed the message - plt.plot(PLOT_DATA) - msg = comm_request({"image/png": None}, msg_type="display_data") - assert hook(msg) is None + assert init_hook(hook) is None # It creates a new figure and comm assert len(hook.figures) == 1 @@ -129,16 +159,6 @@ def test_hook_call(hook: PositronDisplayPublisherHook, images_path: Path) -> Non assert not err -def test_hook_handle_msg_noop_on_unknown_method(figure_comm: DummyComm) -> None: - # Handle a message with an invalid msg_type - msg = json_rpc_request("not_render", {}) - figure_comm.handle_msg(msg) - - assert figure_comm.messages == [ - json_rpc_error(JsonRpcErrorCode.METHOD_NOT_FOUND, "Unknown method 'not_render'") - ] - - def render_request(comm_id: str, width_px: int = 500, height_px: int = 500, pixel_ratio: int = 1): return json_rpc_request( "render", diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_positron_ipkernel.py b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_positron_ipkernel.py index a47cd53aa20..b5b3e89617c 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_positron_ipkernel.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_positron_ipkernel.py @@ -10,8 +10,9 @@ import pandas as pd import pytest from IPython.utils.syspathcontext import prepended_to_syspath -from positron_ipykernel.positron_ipkernel import SessionMode + from positron_ipykernel.help import help +from positron_ipykernel.session_mode import SessionMode from positron_ipykernel.utils import alias_home from .conftest import PositronShell @@ -86,7 +87,7 @@ def test_console_traceback( shell: PositronShell, tmp_path: Path, mock_displayhook: Mock, monkeypatch ) -> None: # Ensure that we're in console mode. - monkeypatch.setattr(shell, "session_mode", SessionMode.Console) + monkeypatch.setattr(shell, "session_mode", SessionMode.CONSOLE) # We follow the approach of IPython's test_ultratb.py, which is to create a temporary module, # prepend its parent directory to sys.path, import it, then run a cell that calls a function @@ -177,7 +178,7 @@ def test_notebook_traceback( shell: PositronShell, tmp_path: Path, mock_displayhook: Mock, monkeypatch ) -> None: # Ensure that we're in notebook mode. - monkeypatch.setattr(shell, "session_mode", SessionMode.Notebook) + monkeypatch.setattr(shell, "session_mode", SessionMode.NOTEBOOK) # We follow the approach of IPython's test_ultratb.py, which is to create a temporary module, # prepend its parent directory to sys.path, import it, then run a cell that calls a function diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/widget.py b/extensions/positron-python/python_files/positron/positron_ipykernel/widget.py index f9682e1cde9..4e6879c24b8 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/widget.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/widget.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) -_WIDGET_MIME_TYPE = "application/vnd.jupyter.widget-view+json" +WIDGET_MIME_TYPE = "application/vnd.jupyter.widget-view+json" @enum.unique @@ -48,11 +48,11 @@ def __call__(self, msg, *args, **kwargs) -> Optional[dict]: if msg["msg_type"] == "display_data": # If there is no widget, let the parent deal with the msg. data = msg["content"]["data"] - if _WIDGET_MIME_TYPE not in data: + if WIDGET_MIME_TYPE not in data: logger.warning("No widget MIME type found.") return msg - comm_id = data[_WIDGET_MIME_TYPE].get("model_id") + comm_id = data[WIDGET_MIME_TYPE].get("model_id") if comm_id is None: logger.warning("No comm associated with widget.") diff --git a/extensions/positron-python/python_files/positron/positron_language_server.py b/extensions/positron-python/python_files/positron/positron_language_server.py index ed1f304f712..4891d8cd36d 100644 --- a/extensions/positron-python/python_files/positron/positron_language_server.py +++ b/extensions/positron-python/python_files/positron/positron_language_server.py @@ -9,8 +9,9 @@ import os import sys -from positron_ipykernel.positron_ipkernel import PositronIPKernelApp, SessionMode +from positron_ipykernel.positron_ipkernel import PositronIPKernelApp from positron_ipykernel.positron_jedilsp import POSITRON +from positron_ipykernel.session_mode import SessionMode logger = logging.getLogger(__name__) @@ -57,7 +58,7 @@ def parse_args() -> argparse.Namespace: "--session-mode", help="session mode in which the kernel is to be started", type=SessionMode, - default=SessionMode.Default, + default=SessionMode.DEFAULT, choices=sorted(SessionMode), ) args = parser.parse_args() diff --git a/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts b/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts index 0c1a1989e4b..89eb57c9526 100644 --- a/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts +++ b/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts @@ -4,7 +4,7 @@ import { Disposable } from 'vs/base/common/lifecycle'; import { IPositronPlotMetadata, PlotClientInstance } from 'vs/workbench/services/languageRuntime/common/languageRuntimePlotClient'; -import { ILanguageRuntimeMessageOutput, RuntimeOutputKind } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService'; +import { ILanguageRuntimeMessageOutput, LanguageRuntimeSessionMode, RuntimeOutputKind } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService'; import { ILanguageRuntimeSession, IRuntimeSessionService, RuntimeClientType } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService'; import { HTMLFileSystemProvider } from 'vs/platform/files/browser/htmlFileSystemProvider'; import { IFileService } from 'vs/platform/files/common/files'; @@ -446,31 +446,34 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe } })); - // Listen for static plots being emitted, and register each one with - // the plots service. - this._register(session.onDidReceiveRuntimeMessageOutput(async (message) => { - // Check to see if we we already have a plot client for this - // message ID. If so, we don't need to do anything. - if (this.hasPlot(session.sessionId, message.id)) { - return; - } + // Configure console-specific behavior. + if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) { + // Listen for static plots being emitted, and register each one with + // the plots service. + this._register(session.onDidReceiveRuntimeMessageOutput(async (message) => { + // Check to see if we we already have a plot client for this + // message ID. If so, we don't need to do anything. + if (this.hasPlot(session.sessionId, message.id)) { + return; + } - const code = this._recentExecutions.has(message.parent_id) ? - this._recentExecutions.get(message.parent_id) : ''; - if (message.kind === RuntimeOutputKind.StaticImage) { - // Create a new static plot client instance and register it with the service. - this.registerStaticPlot(session.sessionId, message, code); + const code = this._recentExecutions.has(message.parent_id) ? + this._recentExecutions.get(message.parent_id) : ''; + if (message.kind === RuntimeOutputKind.StaticImage) { + // Create a new static plot client instance and register it with the service. + this.registerStaticPlot(session.sessionId, message, code); - // Raise the Plots pane so the plot is visible. - this._viewsService.openView(POSITRON_PLOTS_VIEW_ID, false); - } else if (message.kind === RuntimeOutputKind.PlotWidget) { - // Create a new webview plot client instance and register it with the service. - await this.registerWebviewPlot(session, message, code); + // Raise the Plots pane so the plot is visible. + this._viewsService.openView(POSITRON_PLOTS_VIEW_ID, false); + } else if (message.kind === RuntimeOutputKind.PlotWidget) { + // Create a new webview plot client instance and register it with the service. + await this.registerWebviewPlot(session, message, code); - // Raise the Plots pane so the plot is visible. - this._viewsService.openView(POSITRON_PLOTS_VIEW_ID, false); - } - })); + // Raise the Plots pane so the plot is visible. + this._viewsService.openView(POSITRON_PLOTS_VIEW_ID, false); + } + })); + } } /**