diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/matplotlib_backend.py b/extensions/positron-python/python_files/positron/positron_ipykernel/matplotlib_backend.py new file mode 100644 index 00000000000..b1501dea304 --- /dev/null +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/matplotlib_backend.py @@ -0,0 +1,195 @@ +# +# Copyright (C) 2024 Posit Software, PBC. All rights reserved. +# +""" +The matplotlib backend for Positron. +""" + +from __future__ import annotations + +import hashlib +import io +import logging +from typing import Optional, Union, cast + +import matplotlib +from matplotlib._pylab_helpers import Gcf +from matplotlib.backend_bases import FigureManagerBase +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure + +from .plots import Plot + +logger = logging.getLogger(__name__) + + +class FigureManagerPositron(FigureManagerBase): + """ + Interface for the matplotlib backend to interact with the Positron frontend. + + Parameters: + ----------- + canvas + The canvas for this figure. + num + The figure number. + + Attributes: + ----------- + canvas + The canvas for this figure. + """ + + canvas: FigureCanvasPositron + + def __init__(self, canvas: FigureCanvasPositron, num: Union[int, str]): + from .positron_ipkernel import PositronIPyKernel + + super().__init__(canvas, num) + + self._plot: Optional[Plot] = None + self._plots_service = cast(PositronIPyKernel, PositronIPyKernel.instance()).plots_service + + @property + def is_visible(self) -> bool: + """ + Whether the figure is visible to the frontend. + """ + return self._plot is not None + + def show(self) -> None: + """ + Called by matplotlib when a figure is shown via `plt.show()` or `figure.show()`. + """ + if self._plot is None: + # The frontend should respond with a render request, so there's no need for the explicit + # show call in this case. + self._plot = self._plots_service.create_plot(self.canvas._render, self._handle_close) + else: + self._plot.show() + + def destroy(self) -> None: + """ + Called by matplotlib when a figure is closed via `plt.close()`. + """ + if self._plot is not None: + self._plots_service.close_plot(self._plot) + + def update(self) -> None: + """ + Notify the frontend that the plot needs to be rerendered. + + Called by the canvas when a figure is drawn and its contents have changed. + """ + if self._plot is None: + logger.warning("Cannot update a plot that is not visible") + else: + self._plot.update() + + def _handle_close(self) -> None: + """ + Called by the plots service after the plot is closed in the frontend. + """ + # Notify matplotlib to close the figure (and its manager and canvas). + Gcf.destroy(self) + + +class FigureCanvasPositron(FigureCanvasAgg): + """ + The canvas for a figure in the Positron backend. + + Parameters: + ----------- + figure + The figure to draw on this canvas. + + Attributes: + ----------- + manager + The manager for this canvas. + """ + + manager: FigureManagerPositron + + manager_class = FigureManagerPositron # type: ignore + + def __init__(self, figure: Optional[Figure] = None) -> None: + super().__init__(figure) + + # Track the hash of the canvas contents for change detection. + self._last_hash = "" + + def draw(self, is_rendering=False) -> None: + """ + Draw the canvas; send an update event if the canvas has changed. + + Parameters: + ----------- + is_rendering + Whether the canvas is being rendered, to avoid recursively requesting an update from the + frontend. + """ + logger.debug("Drawing to canvas") + try: + super().draw() + finally: + if self.manager.is_visible and not is_rendering: + current_hash = self._hash_buffer_rgba() + logger.debug(f"Canvas: last hash: {self._last_hash[:6]}") + logger.debug(f"Canvas: current hash: {current_hash[:6]}") + if current_hash == self._last_hash: + logger.debug("Canvas: hash is the same, no need to render") + else: + logger.debug("Canvas: hash changed, requesting a render") + self.manager.update() + + def _render(self, width_px: int, height_px: int, pixel_ratio: float, format: str) -> bytes: + # Set the device pixel ratio to the requested value. + self._set_device_pixel_ratio(pixel_ratio) # type: ignore + + # This must be set before setting the size and can't be passed via print_figure else the + # resulting size won't match the request size. + self.figure.set_layout_engine("tight") + + # Resize the figure to the requested size in pixels. + width_in = width_px * self.device_pixel_ratio / self.figure.dpi + height_in = height_px * self.device_pixel_ratio / self.figure.dpi + self.figure.set_size_inches(width_in, height_in, forward=False) + + # Render the canvas. + figure_buffer = io.BytesIO() + with io.BytesIO() as figure_buffer: + self.print_figure( + figure_buffer, + format=format, + dpi=self.figure.dpi, + ) + rendered = figure_buffer.getvalue() + + # NOTE: For some reason, setting the layout engine earlier then calling print_figure + # requires this redraw before calculating the hash else the next draw() call will + # spuriously detect a change. + self.draw(is_rendering=True) + self._last_hash = self._hash_buffer_rgba() + + return rendered + + def _hash_buffer_rgba(self) -> str: + """Hash the canvas contents for change detection.""" + return hashlib.sha1(self.buffer_rgba()).hexdigest() + + +def enable_positron_matplotlib_backend() -> None: + """ + Enable this backend. + """ + # Enable interactive mode to allow for redraws after each cell execution. + matplotlib.interactive(True) + + # Set the backend. + matplotlib.use("module://positron_ipykernel.matplotlib_backend") + + +# Fulfill the matplotlib backend API. +FigureCanvas = FigureCanvasPositron +FigureManager = FigureManagerPositron diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/plot_comm.py b/extensions/positron-python/python_files/positron/positron_ipykernel/plot_comm.py index 57a160b890e..6d8e245777c 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/plot_comm.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/plot_comm.py @@ -98,6 +98,9 @@ class PlotFrontendEvent(str, enum.Enum): # Notification that a plot has been updated on the backend. Update = "update" + # Show a plot. + Show = "show" + PlotResult.update_forward_refs() 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 03b061db47f..d967f3db397 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/plots.py @@ -2,30 +2,21 @@ # Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved. # +from __future__ import annotations + import base64 -import codecs -import io import logging -import pickle import uuid -from typing import Any, Dict, List, Optional - -import comm +from typing import Callable, List, Optional, Protocol -from .plot_comm import PlotBackendMessageContent, PlotResult, RenderRequest -from .positron_comm import CommMessage, JsonRpcErrorCode, PositronComm +from .plot_comm import PlotBackendMessageContent, PlotFrontendEvent, PlotResult, RenderRequest +from .positron_comm import CommMessage, PositronComm from .session_mode import SessionMode from .utils import JsonRecord -from .widget import WIDGET_MIME_TYPE logger = logging.getLogger(__name__) -# Matplotlib Default Figure Size -DEFAULT_WIDTH_IN = 6.4 -DEFAULT_HEIGHT_IN = 4.8 -BASE_DPI = 100 - MIME_TYPE = { "png": "image/png", "svg": "image/svg+xml", @@ -34,218 +25,148 @@ } -class PositronDisplayPublisherHook: - 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.fignums: List[int] = [] - - 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 - - # 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: +class Plot: + """ + The backend representation of a frontend plot instance. + + Paramaters: + ----------- + comm + The communication channel to the frontend plot instance. + render + A callable that renders the plot. See `plot_comm.RenderRequest` for parameter details. + close_callback + A callback that is called after the plot comm is closed. + """ + + def __init__( + self, + comm: PositronComm, + render: Renderer, + close_callback: Optional[Callable[[], None]] = None, + ) -> None: + self._comm = comm + self._render = render + self._close_callback = close_callback + + self._closed = False + + self._comm.on_msg(self._handle_msg, PlotBackendMessageContent) + + @property + def closed(self) -> bool: """ - Create a new plot comm with the given id. + Whether the plot is closed. """ - plot_comm = PositronComm(comm.create_comm(target_name=self.target_name, comm_id=comm_id)) - self.comms[comm_id] = plot_comm - plot_comm.on_msg(self.handle_msg, PlotBackendMessageContent) + return self._closed - def handle_msg(self, msg: CommMessage[PlotBackendMessageContent], raw_msg: JsonRecord) -> None: + def close(self) -> None: """ - Handle client messages to render a plot figure. + Close the plot. """ - comm_id = msg.content.comm_id - request = msg.content.data - - figure_comm = self.comms.get(comm_id, None) - if figure_comm is None: - logger.warning(f"Plot figure comm {comm_id} not found") + if self._closed: return - if isinstance(request, RenderRequest): - pickled = self.figures.get(comm_id, None) - if pickled is None: - figure_comm.send_error( - code=JsonRpcErrorCode.INVALID_PARAMS, message=f"Figure {comm_id} not found" - ) - return - - width_px = request.params.width or 0 - height_px = request.params.height or 0 - pixel_ratio = request.params.pixel_ratio or 1.0 - format = request.params.format or "png" - - if width_px != 0 and height_px != 0: - format_dict = self._resize_pickled_figure( - pickled, width_px, height_px, pixel_ratio, [format] - ) - mime_type = MIME_TYPE[format] - data = format_dict[mime_type] - output = PlotResult(data=data, mime_type=mime_type).dict() - figure_comm.send_result(data=output, metadata={"mime_type": mime_type}) + self._closed = True + self._comm.close() + if self._close_callback: + self._close_callback() - else: - logger.warning(f"Unhandled request: {request}") + def show(self) -> None: + """ + Show the plot. + """ + self._comm.send_event(PlotFrontendEvent.Show, {}) - def shutdown(self) -> None: + def update(self) -> None: """ - Shutdown plot comms and release any resources. + Notify the frontend that the plot needs to be rerendered. """ - for figure_comm in self.comms.values(): - try: - figure_comm.close() - except Exception: - pass - self.comms.clear() - self.figures.clear() + self._comm.send_event(PlotFrontendEvent.Update, {}) - # -- Private Methods -- + def _handle_msg(self, msg: CommMessage[PlotBackendMessageContent], raw_msg: JsonRecord) -> None: + request = msg.content.data + if isinstance(request, RenderRequest): + self._handle_render( + request.params.width, + request.params.height, + request.params.pixel_ratio, + request.params.format, + ) + else: + logger.warning(f"Unhandled request: {request}") - def _pickle_current_figure(self) -> Optional[str]: - # 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 + def _handle_render( + self, + width_px: int, + height_px: int, + pixel_ratio: float, + format: str, + ) -> None: + rendered = self._render(width_px, height_px, pixel_ratio, format) + data = base64.b64encode(rendered).decode() + result = PlotResult(data=data, mime_type=MIME_TYPE[format]).dict() + self._comm.send_result(data=result) - # We turn off interactive mode before accessing the plot context - 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 +class Renderer(Protocol): + """ + A callable that renders a plot. See `plot_comm.RenderRequest` for parameter details. + """ - # Get the current figure, remove it from displayed in the next call - figure = plt.figure(self.fignums.pop(0)) + def __call__(self, width_px: int, height_px: int, pixel_ratio: float, format: str) -> bytes: ... - if self._is_figure_empty(figure): - logger.warning("Figure is empty") - return None - # Pickle the current figure - pickled = codecs.encode(pickle.dumps(figure), "base64").decode() +class PlotsService: + """ + The plots service is responsible for managing `Plot` instances. - return pickled + Paramaters: + ----------- + target_name + The name of the target for plot comms, as defined in the frontend. + session_mode + The session mode that the kernel was started in. + """ - def _resize_pickled_figure( - self, - pickled: str, - new_width_px: int = 614, - new_height_px: int = 460, - pixel_ratio: float = 1.0, - formats: list = ["png"], - ) -> dict: - # 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 - - # Turn off interactive mode before, including before unpickling a - # figures (otherwise it will cause and endless loop of plot changes) - was_interactive = plt.isinteractive() - plt.ioff() - - figure = pickle.loads(codecs.decode(pickled.encode(), "base64")) - figure_buffer = io.BytesIO() - - # Adjust the DPI based on pixel_ratio to accommodate high - # resolution displays... - dpi = BASE_DPI * pixel_ratio - figure.set_dpi(dpi) - figure.set_layout_engine("tight") # eliminates whitespace around the figure - - # ... but use base DPI to convert to inch based dimensions. - width_in, height_in = figure.get_size_inches() - new_width_in = new_width_px / BASE_DPI - new_height_in = new_height_px / BASE_DPI - - # Try to determine if the figure had an explicit width or height set. - if width_in == DEFAULT_WIDTH_IN and height_in == DEFAULT_HEIGHT_IN: - # If default values are still set, apply new size, even if this - # resets the aspect ratio - width_in = new_width_in - height_in = new_height_in - else: - # Preserve the existing aspect ratio, constraining the scale - # based on the shorter dimension - if width_in < height_in: - height_in = height_in * (new_width_in / width_in) - width_in = new_width_in - else: - width_in = width_in * (new_height_in / height_in) - height_in = new_height_in - - figure.set_size_inches(width_in, height_in) - - # Render the figure to a buffer - # using format_display_data() crops the figure to smaller than requested size - figure.savefig(figure_buffer, format=formats[0]) - figure_buffer.seek(0) - image_data = base64.b64encode(figure_buffer.read()).decode() - key = MIME_TYPE[formats[0]] - - format_dict = {key: image_data} - - plt.close(figure) - - if was_interactive: - plt.ion() - return format_dict - - def _is_figure_empty(self, figure): - children = figure.get_children() - if len(children) < 1: - return True - - for child in children: - if child.get_visible(): - return False - - return True + def __init__(self, target_name: str, session_mode: SessionMode): + self._target_name = target_name + self._session_mode = session_mode + + self._plots: List[Plot] = [] + + def create_plot(self, render: Renderer, close_callback: Callable[[], None]) -> Plot: + """ + Create a plot. + + See Also: + --------- + Plot + """ + comm_id = str(uuid.uuid4()) + logger.info(f"Creating plot with comm {comm_id}") + plot_comm = PositronComm.create(self._target_name, comm_id) + plot = Plot(plot_comm, render, close_callback) + self._plots.append(plot) + return plot + + def close_plot(self, plot: Plot) -> None: + """ + Close a plot. + + Parameters: + ----------- + plot + The plot to close. + """ + if plot.closed: + return + plot.close() + self._plots.remove(plot) + + def shutdown(self) -> None: + """ + Shutdown the plots service. + """ + for plot in list(self._plots): + self.close_plot(plot) 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 d534430f7bc..4c8cfb06370 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 @@ -60,15 +60,70 @@ class JsonRpcErrorCode(enum.IntEnum): class CommMessage(GenericModel, Generic[T_content]): + """ + A generic message received from the frontend-side of a comm. + """ + content: T_content class PositronComm: - """A wrapper around a base IPython comm that provides a JSON-RPC interface""" + """ + A wrapper around a base IPython comm that provides a JSON-RPC interface. + + Paramaters: + ----------- + comm + The wrapped IPython comm. + + Attributes: + ----------- + comm + The wrapped IPython comm. + comm_id + """ def __init__(self, comm: comm.base_comm.BaseComm) -> None: self.comm = comm + @classmethod + def create(cls, target_name: str, comm_id: str) -> PositronComm: + """ + Create a Positron comm. + + Parameters: + ----------- + target_name + The name of the target for the comm, as defined in the frontend. + comm_id + The unique identifier for the comm. + + Returns: + -------- + PositronComm + The new PositronComm instance. + """ + base_comm = comm.create_comm(target_name=target_name, comm_id=comm_id) + return cls(base_comm) + + @property + def comm_id(self) -> str: + """ + The unique identifier of this comm. + """ + return self.comm.comm_id + + def on_close(self, callback: Callable[[JsonRecord], None]): + """ + Register a callback for when the frontend-side version of this comm is closed. + + Paramaters: + ----------- + callback + Called when the comm is closed, with the raw close message. + """ + self.comm.on_close(callback) + def on_msg( self, callback: Callable[[CommMessage[T_content], JsonRecord], None], @@ -77,9 +132,14 @@ def on_msg( """ Register a callback for an RPC request from the frontend. - Will be called with both the parsed `msg: CommMessage` and the original `raw_msg`. - - If the `raw_msg` could not be parsed, a JSON-RPC error will be sent to the frontend. + Parameters: + ----------- + callback + Called when a message is received, with both the parsed message `msg: CommMessage` and + original `raw_msg`. Not called if the `raw_msg` could not be parsed; instead, a JSON-RPC + error will be sent to the frontend. + content_cls + The Pydantic model to parse the message with. """ def handle_msg( @@ -128,7 +188,16 @@ def handle_msg( self.comm.on_msg(handle_msg) def send_result(self, data: JsonData = None, metadata: Optional[JsonRecord] = None) -> None: - """Send a JSON-RPC result to the frontend-side version of this comm""" + """ + Send a JSON-RPC result to the frontend-side version of this comm. + + Parameters: + ----------- + data + The result data to send. + metadata + The metadata to send with the result. + """ result = dict( jsonrpc="2.0", result=data, @@ -140,7 +209,16 @@ def send_result(self, data: JsonData = None, metadata: Optional[JsonRecord] = No ) def send_event(self, name: str, payload: JsonRecord) -> None: - """Send a JSON-RPC notification (event) to the frontend-side version of this comm""" + """ + Send a JSON-RPC notification (event) to the frontend-side version of this comm. + + Parameters: + ----------- + name + The name of the event. + payload + The payload of the event. + """ event = dict( jsonrpc="2.0", method=name, @@ -149,7 +227,16 @@ def send_event(self, name: str, payload: JsonRecord) -> None: self.comm.send(data=event) def send_error(self, code: JsonRpcErrorCode, message: Optional[str] = None) -> None: - """Send a JSON-RPC result to the frontend-side version of this comm""" + """ + Send a JSON-RPC result to the frontend-side version of this comm. + + Parameters: + ----------- + code + The error code to send. + message + The error message to send. + """ error = dict( jsonrpc="2.0", error=dict( @@ -164,5 +251,7 @@ def send_error(self, code: JsonRpcErrorCode, message: Optional[str] = None) -> N ) def close(self) -> None: - """Close the underlying comm.""" + """ + Close the frontend-side version of this comm. + """ self.comm.close() 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 c6b1d6e3653..c73e756d108 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,20 +20,15 @@ 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 from .data_explorer import DataExplorerService from .help import HelpService, help from .lsp import LSPService -from .plots import PositronDisplayPublisherHook +from .matplotlib_backend import enable_positron_matplotlib_backend +from .plots import PlotsService from .session_mode import SessionMode from .ui import UiService from .utils import JsonRecord @@ -224,6 +219,7 @@ def _handle_post_run_cell(self, info: ExecutionInfo) -> None: After execution, sends an update message to the client to summarize the changes observed to variables in the user's environment. """ + # TODO: Split these to separate callbacks? # Check for changes to the working directory try: self.kernel.ui_service.poll_working_directory() @@ -343,7 +339,7 @@ def __init__(self, **kwargs) -> None: # Create Positron services self.data_explorer_service = DataExplorerService(_CommTarget.DataExplorer) - self.display_pub_hook = PositronDisplayPublisherHook(_CommTarget.Plot, self.session_mode) + self.plots_service = PlotsService(_CommTarget.Plot, self.session_mode) self.ui_service = UiService() self.help_service = HelpService() self.lsp_service = LSPService(self) @@ -359,7 +355,6 @@ def __init__(self, **kwargs) -> None: _CommTarget.Variables, self.variables_service.on_comm_open ) # Register display publisher hooks - self.shell.display_pub.register_hook(self.display_pub_hook) self.shell.display_pub.register_hook(self.widget_hook) # Ignore warnings that the user can't do anything about @@ -398,10 +393,10 @@ async def do_shutdown(self, restart: bool) -> JsonRecord: # type: ignore Report # Shutdown Positron services self.data_explorer_service.shutdown() - self.display_pub_hook.shutdown() self.ui_service.shutdown() self.help_service.shutdown() self.lsp_service.shutdown() + self.plots_service.shutdown() self.widget_hook.shutdown() await self.variables_service.shutdown() self.connections_service.shutdown() @@ -413,12 +408,25 @@ async def do_shutdown(self, restart: bool) -> JsonRecord: # type: ignore Report class PositronIPKernelApp(IPKernelApp): + kernel: PositronIPyKernel + # Use the PositronIPyKernel class. kernel_class: Type[PositronIPyKernel] = traitlets.Type(PositronIPyKernel) # type: ignore # Positron-specific attributes: session_mode: SessionMode = SessionMode.trait() # type: ignore + def init_gui_pylab(self): + try: + # Enable the Positron matplotlib backend if we're not in a notebook. + # If we're in a notebook, use IPython's default backend via the super() call below. + if self.session_mode != SessionMode.NOTEBOOK: + enable_positron_matplotlib_backend() + except Exception: + logger.error("Error setting matplotlib backend") + + return super().init_gui_pylab() + # # OSC8 functionality 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 e82cf1a506c..a3c016f60d9 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 @@ -33,6 +33,11 @@ def publish_msg(self, msg_type, **msg): # type: ignore ReportIncompatibleMethod msg["msg_type"] = msg_type self.messages.append(msg) + def pop_messages(self): + messages = list(self.messages) + self.messages.clear() + return messages + # Enable autouse so that all comms are created as DummyComms. @pytest.fixture(autouse=True) 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 6462d49cdb5..e0253f8fa2e 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 @@ -3,51 +3,61 @@ # import base64 -import codecs import io -import pickle from pathlib import Path -from typing import Iterable, Optional, cast +from typing import Iterable, List, cast import matplotlib import matplotlib.pyplot as plt import pytest -from IPython.core.formatters import DisplayFormatter -from matplotlib.axes import Axes -from matplotlib.figure import Figure -from matplotlib.testing.compare import compare_images -from matplotlib_inline.backend_inline import configure_inline_support +from PIL import Image -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 positron_ipykernel.matplotlib_backend import enable_positron_matplotlib_backend +from positron_ipykernel.plots import PlotsService +from positron_ipykernel.positron_ipkernel import PositronIPyKernel, _CommTarget from .conftest import DummyComm, PositronShell -from .utils import comm_open_message, comm_request, json_rpc_error, json_rpc_request +from .utils import ( + comm_close_message, + comm_open_message, + json_rpc_notification, + json_rpc_request, + json_rpc_response, +) + +# +# Matplotlib backend + shell + plots service integration tests. +# -PLOT_DATA = [1, 2] TARGET_NAME = "target_name" @pytest.fixture(autouse=True) -def setup_matplotlib(shell: PositronShell) -> Iterable[None]: - # Use IPython's `matplotlib_inline` backend - backend = "module://matplotlib_inline.backend_inline" - matplotlib.use(backend) +def setup_positron_matplotlib_backend() -> None: + enable_positron_matplotlib_backend() + + assert matplotlib.get_backend() == "module://positron_ipykernel.matplotlib_backend" + assert matplotlib.is_interactive() + - # Enable all IPython mimetype formatters - display_formatter = cast(DisplayFormatter, shell.display_formatter) - active_types = display_formatter.active_types - display_formatter.active_types = display_formatter.format_types +@pytest.fixture(autouse=True) +def import_pyplot(shell: PositronShell) -> None: + # Import pyplot for convenience. + shell.run_cell("import matplotlib.pyplot as plt") + + +@pytest.fixture +def plots_service(kernel: PositronIPyKernel) -> Iterable[PlotsService]: + """ + The Positron plots service. + """ + plots_service = kernel.plots_service - # Enable matplotlib IPython formatters - configure_inline_support(shell, backend) + assert not plots_service._plots - yield + yield plots_service - # Restore the original active formatters - display_formatter.active_types = active_types + plt.close("all") @pytest.fixture(scope="session") @@ -57,206 +67,172 @@ def images_path() -> Path: return images_path -@pytest.fixture -def hook() -> PositronDisplayPublisherHook: - return PositronDisplayPublisherHook(TARGET_NAME, SessionMode.CONSOLE) +def test_mpl_dont_create_plot_on_new_figure( + shell: PositronShell, plots_service: PlotsService +) -> None: + # Creating a figure should not yet create a plot with the plots service. + shell.run_cell("plt.figure()") + assert not plots_service._plots -@pytest.fixture -def notebook_hook() -> PositronDisplayPublisherHook: - return PositronDisplayPublisherHook(TARGET_NAME, SessionMode.NOTEBOOK) +def _get_plot_comms(plots_service: PlotsService) -> List[DummyComm]: + return [cast(DummyComm, plot._comm.comm) for plot in plots_service._plots] -def display_data_message() -> JsonRecord: - """ - 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 _get_single_plot_comm(plots_service: PlotsService) -> DummyComm: + plot_comms = _get_plot_comms(plots_service) + assert len(plot_comms) == 1 + return plot_comms[0] -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) - try: - msg = display_data_message() - return hook(msg) - finally: - plt.close() +@pytest.mark.parametrize("code", ["plt.figure(); plt.show()", "plt.figure().show()"]) +def test_mpl_send_open_comm_on_plt_show( + code: str, shell: PositronShell, plots_service: PlotsService +) -> None: + # Showing a figure should create a plot with the plots service and open a corresponding comm. + shell.run_cell(code) + plot_comm = _get_single_plot_comm(plots_service) + assert plot_comm.pop_messages() == [comm_open_message(_CommTarget.Plot)] -@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 +def _create_mpl_plot(shell: PositronShell, plots_service: PlotsService) -> DummyComm: + shell.run_cell("plt.figure().show()") + plot_comm = _get_single_plot_comm(plots_service) + plot_comm.messages.clear() + return plot_comm - # 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)] +def test_mpl_send_show_on_successive_plt_show( + shell: PositronShell, plots_service: PlotsService +) -> None: + plot_comm = _create_mpl_plot(shell, plots_service) - # Clear messages due to the comm_open - figure_comm.messages.clear() + # Show the figure again. + shell.run_cell("plt.show()") + assert plot_comm.pop_messages() == [json_rpc_notification("show", {})] - return figure_comm + # It should also work with Figure.show(). + shell.run_cell("plt.gcf().show()") + assert plot_comm.pop_messages() == [json_rpc_notification("show", {})] -def test_hook_call_noop_on_non_display_data(hook: PositronDisplayPublisherHook) -> None: - msg = comm_request({"image/png": None}, msg_type="not_display_data") - assert hook(msg) == msg - assert hook.figures == {} - assert hook.comms == {} +def test_mpl_send_update_on_draw(shell: PositronShell, plots_service: PlotsService) -> None: + plot_comm = _create_mpl_plot(shell, plots_service) + # Drawing to an active plot should trigger an update. + shell.run_cell("plt.plot([1, 2])") + assert plot_comm.pop_messages() == [json_rpc_notification("update", {})] -def test_hook_call_noop_on_no_image_png(hook: PositronDisplayPublisherHook) -> None: - msg = comm_request({}, msg_type="display_data") - assert hook(msg) == msg - assert hook.figures == {} - assert hook.comms == {} +def test_mpl_dont_send_update_on_execution( + shell: PositronShell, plots_service: PlotsService +) -> None: + plot_comm = _create_mpl_plot(shell, plots_service) -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 == {} + # Executing code that doesn't draw to the active plot should not trigger an update. + shell.run_cell("1") + assert plot_comm.pop_messages() == [] -def test_hook_call(hook: PositronDisplayPublisherHook, images_path: Path) -> None: - assert init_hook(hook) is None +def test_mpl_send_close_comm_on_plt_close( + shell: PositronShell, plots_service: PlotsService +) -> None: + plot_comm = _create_mpl_plot(shell, plots_service) - # It creates a new figure and comm - assert len(hook.figures) == 1 - id = next(iter(hook.figures)) - assert id in hook.comms + # Closing the figure should close the plot and send a comm close message. + shell.run_cell("plt.close()") + assert plot_comm.pop_messages() == [comm_close_message()] + assert not plots_service._plots - # Check the comm's properties - comm = hook.comms[id].comm - assert comm.target_name == hook.target_name - assert comm.comm_id == id - # Check that the figure is a pickled base64-encoded string by decoding it and comparing it - # with a reference figure. - # First, save the hook's figure - fig_encoded = hook.figures[id] - fig: Figure = pickle.loads(codecs.decode(fig_encoded.encode(), "base64")) - actual = images_path / "test-hook-call-actual.png" - fig.savefig(str(actual)) +def test_mpl_multiple_figures(shell: PositronShell, plots_service: PlotsService) -> None: + # Create two figures, and show them both. + shell.run_cell("f1 = plt.figure(); f2 = plt.figure(); plt.show()") - # Create the reference figure - fig_ref = cast(Figure, plt.figure()) - fig_axes = cast(Axes, fig_ref.subplots()) - fig_axes.plot(PLOT_DATA) - expected = images_path / "test-hook-call-expected.png" - fig_ref.savefig(str(expected)) + plot_comms = _get_plot_comms(plots_service) + assert len(plot_comms) == 2 + for plot_comm in plot_comms: + assert plot_comm.pop_messages() == [comm_open_message(_CommTarget.Plot)] - # Compare actual versus expected figures - err = compare_images(str(actual), str(expected), tol=0) - assert not err + # Draw to the first figure. + shell.run_cell("plt.figure(f1); plt.plot([1, 2])") + assert plot_comms[0].pop_messages() == [json_rpc_notification("update", {})] + assert plot_comms[1].pop_messages() == [] -def render_request(comm_id: str, width_px: int = 500, height_px: int = 500, pixel_ratio: int = 1): - return json_rpc_request( - "render", - {"width": width_px, "height": height_px, "pixel_ratio": pixel_ratio, "format": "png"}, - comm_id=comm_id, - ) + # Draw to the second figure. + shell.run_cell("plt.figure(f2); plt.plot([1, 2])") + assert plot_comms[0].pop_messages() == [] + assert plot_comms[1].pop_messages() == [json_rpc_notification("update", {})] -def test_hook_render_noop_on_unknown_comm(figure_comm: DummyComm) -> None: - # Handle a valid message but invalid comm_id - msg = render_request("unknown_comm_id") - figure_comm.handle_msg(msg) + # Show the first figure. + shell.run_cell("f1.show()") - # No messages sent - assert figure_comm.messages == [] + assert plot_comms[0].pop_messages() == [json_rpc_notification("show", {})] + assert plot_comms[1].pop_messages() == [] + # Show the second figure. + shell.run_cell("f2.show()") -def test_hook_render_error_on_unknown_figure( - hook: PositronDisplayPublisherHook, figure_comm: DummyComm -) -> None: - # Clear the hook's figures to simulate a missing figure - hook.figures.clear() + assert plot_comms[0].pop_messages() == [] + assert plot_comms[1].pop_messages() == [json_rpc_notification("show", {})] - # Handle a message with a valid msg_type and valid comm_id, but the hook now has a missing figure - msg = render_request(figure_comm.comm_id) - figure_comm.handle_msg(msg) - # Check that we receive an error reply - assert figure_comm.messages == [ - json_rpc_error(JsonRpcErrorCode.INVALID_PARAMS, f"Figure {figure_comm.comm_id} not found") +def test_mpl_render(shell: PositronShell, plots_service: PlotsService, images_path: Path) -> None: + # First show the figure and get the plot comm. + shell.run_cell("plt.plot([1, 2])\nplt.show()") + plot_comm = _get_single_plot_comm(plots_service) + assert plot_comm.pop_messages() == [ + comm_open_message(_CommTarget.Plot), + # NOTE: The update here is unnecessary since when the frontend receives a comm open, it + # responds with a render request. It is probably harmless though since the frontend + # debounces render requests. It happens because all figures are redrawn post cell execution, + # when matplotlib interactive mode is enabled. + json_rpc_notification("update", {}), ] + # Send a render request to the plot comm. + width = 400 + height = 300 + pixel_ratio = 2 + format = "png" + msg = json_rpc_request( + "render", + {"width": width, "height": height, "pixel_ratio": pixel_ratio, "format": format}, + comm_id="dummy_comm_id", + ) + plot_comm.handle_msg(msg) -def _save_base64_image(encoded: str, filename: Path) -> None: - image = codecs.decode(encoded.encode(), "base64") - with open(filename, "wb") as f: - f.write(image) - - -def test_hook_render(figure_comm: DummyComm, images_path: Path) -> None: - # Send a valid render message with a custom width and height - width_px = height_px = 100 - pixel_ratio = 1 - msg = render_request(figure_comm.comm_id, width_px, height_px, pixel_ratio) - figure_comm.handle_msg(msg) - - # Check that the reply is a comm_msg - reply = figure_comm.messages[0] - assert reply["msg_type"] == "comm_msg" - assert reply["buffers"] is None - assert reply["metadata"] == {"mime_type": "image/png"} - - # Check that the reply data is an `image` message - image_msg = reply["data"] - assert image_msg["result"]["mime_type"] == "image/png" - - # Check that the reply data includes the expected base64-encoded resized image - - # Save the reply's image - actual = images_path / "test-hook-render-actual.png" - _save_base64_image(image_msg["result"]["data"], actual) - - # Create the reference figure - dpi = BASE_DPI * pixel_ratio - width_in = width_px / BASE_DPI - height_in = height_px / BASE_DPI + responses = plot_comm.pop_messages() + assert len(responses) == 1 + response = responses[0] - fig_buffer = io.BytesIO() - fig_ref = cast(Figure, plt.figure()) - fig_axes = cast(Axes, fig_ref.subplots()) - fig_axes.plot([1, 2]) - fig_ref.set_dpi(dpi) - fig_ref.set_size_inches(width_in, height_in) - fig_ref.set_layout_engine("tight") + # Check that the response includes the expected base64-encoded resized image. + image_bytes = response["data"]["result"].pop("data") + image = Image.open(io.BytesIO(base64.b64decode(image_bytes))) + assert image.format == format.upper() + assert image.size == (width * pixel_ratio, height * pixel_ratio) + # Save it to disk for manual inspection. + image.save(images_path / "test-mpl-render.png") - # Serialize the reference figure as a base64-encoded image - fig_ref.savefig(fig_buffer, format="png") - fig_buffer.seek(0) - expected = images_path / "test-hook-render-expected.png" - _save_base64_image(base64.b64encode(fig_buffer.read()).decode(), expected) + # Check the rest of the response. + assert response == json_rpc_response({"mime_type": "image/png"}) - # Compare the actual vs expected figures - err = compare_images(str(actual), str(expected), tol=0) - assert not err +def test_mpl_shutdown(shell: PositronShell, plots_service: PlotsService) -> None: + # Create a figure and show it. + shell.run_cell("plt.figure(); plt.figure(); plt.show()") + plot_comms = _get_plot_comms(plots_service) -# It's important that we depend on the figure_comm fixture too, so that the hook is initialized -def test_shutdown(hook: PositronDisplayPublisherHook, figure_comm: DummyComm) -> None: - # Double-check that it still has figures and comms - assert len(hook.figures) == 1 - assert len(hook.comms) == 1 + # Double-check that it still has plots. + assert len(plots_service._plots) == 2 - # Double-check that the comm is not yet closed - assert not figure_comm._closed + # Double-check that all comms are still open. + assert not any(comm._closed for comm in plot_comms) - hook.shutdown() + plots_service.shutdown() - # Figures and comms are closed and cleared - assert not hook.figures - assert not hook.comms - assert figure_comm._closed + # Plots are closed and cleared. + assert not plots_service._plots + assert all(comm._closed for comm in plot_comms) diff --git a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/utils.py b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/utils.py index c24a8457e58..703db0d6afc 100644 --- a/extensions/positron-python/python_files/positron/positron_ipykernel/tests/utils.py +++ b/extensions/positron-python/python_files/positron/positron_ipykernel/tests/utils.py @@ -71,6 +71,13 @@ def comm_open_message(target_name: str, data: Optional[JsonRecord] = None) -> Js } +def comm_close_message() -> JsonRecord: + return { + **comm_message(), + "msg_type": "comm_close", + } + + def json_rpc_error(code: int, message: str) -> JsonRecord: return comm_message( { diff --git a/positron/comms/plot-frontend-openrpc.json b/positron/comms/plot-frontend-openrpc.json index 4a24bb86447..a38bc2cacbb 100644 --- a/positron/comms/plot-frontend-openrpc.json +++ b/positron/comms/plot-frontend-openrpc.json @@ -9,6 +9,11 @@ "name": "update", "summary": "Notification that a plot has been updated on the backend.", "params": [] + }, + { + "name": "show", + "summary": "Show a plot.", + "params": [] } ] } diff --git a/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts b/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts index f2eaa68e6d4..8ddc1d3fe5f 100644 --- a/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts +++ b/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts @@ -507,14 +507,23 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe StorageScope.WORKSPACE); }); - // Raise the plot if it's updated by the runtime - plotClient.onDidRenderUpdate((_plot) => { + const selectPlot = () => { // Raise the Plots pane so the user can see the updated plot this._viewsService.openView(POSITRON_PLOTS_VIEW_ID, false); // Select the plot to bring it into view within the history; it's // possible that it is not the most recently created plot this._onDidSelectPlot.fire(plotClient.id); + }; + + // Raise the plot if it's updated by the runtime + plotClient.onDidRenderUpdate((_plot) => { + selectPlot(); + }); + + // Focus the plot if the runtime requests it + plotClient.onDidShowPlot(() => { + selectPlot(); }); // Dispose the plot client when this service is disposed (we own this diff --git a/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts b/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts index 5b19883716f..2cec7b20f2c 100644 --- a/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts +++ b/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts @@ -195,6 +195,9 @@ export class PlotClientInstance extends Disposable implements IPositronPlotClien onDidRenderUpdate: Event; private readonly _renderUpdateEmitter = new Emitter(); + onDidShowPlot: Event; + private readonly _didShowPlotEmitter = new Emitter(); + /** * Creates a new plot client instance. * @@ -227,6 +230,9 @@ export class PlotClientInstance extends Disposable implements IPositronPlotClien // Connect the render update emitter event this.onDidRenderUpdate = this._renderUpdateEmitter.event; + // Connect the show plot emitter event + this.onDidShowPlot = this._didShowPlotEmitter.event; + // Listen to our own state changes this.onDidChangeState((state) => { this._state = state; @@ -238,6 +244,11 @@ export class PlotClientInstance extends Disposable implements IPositronPlotClien this._renderUpdateEmitter.fire(rendered); }); + // Listn for plot show events + this._comm.onDidShow(async (_evt) => { + this._didShowPlotEmitter.fire(); + }); + // Register the client instance with the runtime, so that when this instance is disposed, // the runtime will also dispose the client. this._register(this._comm); diff --git a/src/vs/workbench/services/languageRuntime/common/positronPlotComm.ts b/src/vs/workbench/services/languageRuntime/common/positronPlotComm.ts index 95836a370f3..46bb1f2ea19 100644 --- a/src/vs/workbench/services/languageRuntime/common/positronPlotComm.ts +++ b/src/vs/workbench/services/languageRuntime/common/positronPlotComm.ts @@ -32,14 +32,22 @@ export interface PlotResult { export interface UpdateEvent { } +/** + * Event: Show a plot. + */ +export interface ShowEvent { +} + export enum PlotFrontendEvent { - Update = 'update' + Update = 'update', + Show = 'show' } export class PositronPlotComm extends PositronBaseComm { constructor(instance: IRuntimeClientInstance) { super(instance); this.onDidUpdate = super.createEventEmitter('update', []); + this.onDidShow = super.createEventEmitter('show', []); } /** @@ -64,5 +72,9 @@ export class PositronPlotComm extends PositronBaseComm { * Notification that a plot has been updated on the backend. */ onDidUpdate: Event; + /** + * Show a plot. + */ + onDidShow: Event; }