Skip to content

Commit

Permalink
render plots inline in notebook sessions
Browse files Browse the repository at this point in the history
Plus some minor general housekeeping.
  • Loading branch information
seeM committed Apr 12, 2024
1 parent 7cbd63a commit abd651e
Show file tree
Hide file tree
Showing 10 changed files with 205 additions and 130 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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:
"""
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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__)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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__)
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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)

Expand Down
Loading

0 comments on commit abd651e

Please sign in to comment.