Merge pull request #2592 from posit-dev/bugfix/plot-crop

Fix plot rendering getting cropped

extensions/positron-python/pythonFiles/positron/positron_ipykernel/plots.py

@@ -2,14 +2,15 @@
 # Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved.
 #
 
+import base64
 import codecs
+import io
 import logging
 import pickle
 import uuid
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
 import comm
-from IPython.core.formatters import format_display_data
 
 from .plot_comm import PlotBackendMessageContent, PlotResult, RenderRequest
 from .positron_comm import CommMessage, JsonRpcErrorCode, PositronComm
@@ -22,7 +23,7 @@
 # Matplotlib Default Figure Size
 DEFAULT_WIDTH_IN = 6.4
 DEFAULT_HEIGHT_IN = 4.8
-BASE_DPI = 96
+BASE_DPI = 100
 
 
 class PositronDisplayPublisherHook:
@@ -96,12 +97,10 @@ def handle_msg(self, msg: CommMessage[PlotBackendMessageContent], raw_msg: JsonR
             pixel_ratio = request.params.pixel_ratio or 1.0
 
             if width_px != 0 and height_px != 0:
-                format_dict, md_dict = self._resize_pickled_figure(
-                    pickled, width_px, height_px, pixel_ratio
-                )
+                format_dict = self._resize_pickled_figure(pickled, width_px, height_px, pixel_ratio)
                 data = format_dict["image/png"]
                 output = PlotResult(data=data, mime_type="image/png").dict()
-                figure_comm.send_result(data=output, metadata=md_dict)
+                figure_comm.send_result(data=output, metadata={"mime_type": "image/png"})
 
         else:
             logger.warning(f"Unhandled request: {request}")
@@ -157,7 +156,7 @@ def _resize_pickled_figure(
         new_height_px: int = 460,
         pixel_ratio: float = 1.0,
         formats: list = ["image/png"],
-    ) -> Tuple[dict, dict]:
+    ) -> 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
@@ -168,11 +167,13 @@ def _resize_pickled_figure(
         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()
@@ -197,14 +198,19 @@ def _resize_pickled_figure(
 
         figure.set_size_inches(width_in, height_in)
 
-        format_dict, md_dict = format_display_data(figure, include=formats, exclude=[])  # type: ignore
+        # Render the figure to a buffer
+        # using format_display_data() crops the figure to smaller than requested size
+        figure.savefig(figure_buffer, format="png")
+        figure_buffer.seek(0)
+        image_data = base64.b64encode(figure_buffer.read()).decode()
+
+        format_dict = {"image/png": image_data}
 
         plt.close(figure)
 
         if was_interactive:
             plt.ion()
-
-        return (format_dict, md_dict)
+        return format_dict
 
     def _is_figure_empty(self, figure):
         children = figure.get_children()

extensions/positron-python/pythonFiles/positron/positron_ipykernel/tests/test_plots.py

@@ -2,15 +2,17 @@
 # Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved.
 #
 
+import base64
 import codecs
+import io
 import pickle
 from pathlib import Path
 from typing import Iterable, cast
 
 import matplotlib
 import matplotlib.pyplot as plt
 import pytest
-from IPython.core.formatters import DisplayFormatter, format_display_data
+from IPython.core.formatters import DisplayFormatter
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from matplotlib.testing.compare import compare_images
@@ -187,7 +189,7 @@ def test_hook_render(figure_comm: DummyComm, images_path: Path) -> None:
     reply = figure_comm.messages[0]
     assert reply["msg_type"] == "comm_msg"
     assert reply["buffers"] is None
-    assert reply["metadata"] == {}
+    assert reply["metadata"] == {"mime_type": "image/png"}
 
     # Check that the reply data is an `image` message
     image_msg = reply["data"]
@@ -204,16 +206,19 @@ def test_hook_render(figure_comm: DummyComm, images_path: Path) -> None:
     width_in = width_px / BASE_DPI
     height_in = height_px / BASE_DPI
 
+    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")
 
     # Serialize the reference figure as a base64-encoded image
-    data_ref, _ = format_display_data(fig_ref, include=["image/png"], exclude=[])  # type: ignore
+    fig_ref.savefig(fig_buffer, format="png")
+    fig_buffer.seek(0)
     expected = images_path / "test-hook-render-expected.png"
-    _save_base64_image(data_ref["image/png"], expected)
+    _save_base64_image(base64.b64encode(fig_buffer.read()).decode(), expected)
 
     # Compare the actual vs expected figures
     err = compare_images(str(actual), str(expected), tol=0)