Render plots inline in notebook sessions

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,

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
 
 

extensions/positron-python/python_files/positron/positron_ipykernel/positron_ipkernel.py

@@ -20,14 +20,21 @@
 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 .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)

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__)

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)