Capture the last_plot within the plotting context

Author: has2k1

plotnine/_utils/context.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+from copy import copy
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import pandas as pd
 from packaging.version import Version
 
+from ..session import set_last_plot
+
 if TYPE_CHECKING:
     from typing_extensions import Self
 
@@ -38,6 +41,7 @@ def is_closed(fig) -> bool:
     return not plt.fignum_exists(fig.number)
 
 
+@dataclass
 class plot_context:
     """
     Context within which the plot is built
@@ -50,14 +54,14 @@ class plot_context:
         Whether to show the plot.
     """
 
-    def __init__(self, plot: ggplot, show: bool = False):
-        import matplotlib as mpl
+    plot: ggplot
+    show: bool = False
 
-        self.plot = plot
-        self.show = show
+    def __post_init__(self):
+        import matplotlib as mpl
 
         # Contexts
-        self.rc_context = mpl.rc_context(plot.theme.rcParams)
+        self.rc_context = mpl.rc_context(self.plot.theme.rcParams)
         if PANDAS_LT_3:
             self.pd_option_context = pd.option_context(
                 "mode.copy_on_write",
@@ -69,6 +73,7 @@ def __enter__(self) -> Self:
         Enclose in matplolib & pandas environments
         """
 
+        self._last_plot = copy(self.plot)
         self.rc_context.__enter__()
         if PANDAS_LT_3:
             self.pd_option_context.__enter__()
@@ -86,6 +91,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
                 plt.show()
             else:
                 plt.close(self.plot.figure)
+            set_last_plot(self._last_plot)
         else:
             # There is an exception, close any figure
             if hasattr(self.plot, "figure"):
@@ -127,6 +133,7 @@ def __enter__(self) -> Self:
         """
         Enclose in matplolib & pandas environments
         """
+        self._last_plot = copy(self.cmp)
         self._rc_context.__enter__()
         return self
 
@@ -140,6 +147,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
                 plt.show()
             else:
                 plt.close(self.cmp.figure)
+            set_last_plot(self._last_plot)
         else:
             # There is an exception, close any figure
             if hasattr(self.cmp, "figure"):

plotnine/composition/_compose.py

@@ -18,7 +18,6 @@
 from ..composition._plot_layout import plot_layout
 from ..composition._types import ComposeAddable
 from ..options import get_option
-from ..session import set_last_plot
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -560,7 +559,6 @@ def _draw(cmp):
             self.theme.apply()
             figure.set_layout_engine(PlotnineLayoutEngine(self))
 
-        set_last_plot(self)
         return figure
 
     def _draw_plots(self):

plotnine/ggplot.py

@@ -42,7 +42,6 @@
 from .mapping.aes import aes
 from .options import get_option
 from .scales.scales import Scales
-from .session import set_last_plot
 from .themes.theme import theme, theme_get
 
 if TYPE_CHECKING:
@@ -376,7 +375,6 @@ def draw(self, *, show: bool = False) -> Figure:
             self.theme.apply()
             figure.set_layout_engine(PlotnineLayoutEngine(self))
 
-        set_last_plot(self)
         return figure
 
     def _setup(self) -> Figure: