Add plot_layout(guides="collect") for composition-wide legend collection

doc/_quartodoc.yml

@@ -549,6 +549,7 @@ quartodoc:
         - plot_spacer
         - plot_layout
         - inset_element
+        - guide_area
 
     - title: Options
       desc: |

plotnine/_mpl/layout_manager/_composition_layout_items.py

@@ -13,6 +13,7 @@
     from matplotlib.patches import Rectangle
 
     from plotnine.composition._compose import Compose
+    from plotnine.iapi import legend_artists
 
     from ._composition_side_space import CompositionSideSpaces
 
@@ -22,7 +23,7 @@ class CompositionLayoutItems:
     plot_annotation artists
     """
 
-    def __init__(self, cmp: Compose):
+    def __init__(self, cmp: Compose, *, is_root: bool = True):
         def get(name: str) -> Any:
             """
             Return themeable target or None
@@ -36,6 +37,7 @@ def get(name: str) -> Any:
                 return t
 
         self.cmp = cmp
+        self.is_root = is_root
         self.geometry = ArtistGeometry(cmp.figure)
 
         self.plot_title: Text | None = get("plot_title")
@@ -46,14 +48,25 @@ def get(name: str) -> Any:
             "plot_footer_background"
         )
         self.plot_footer_line: Line2D | None = get("plot_footer_line")
+        self.legends: legend_artists | None = get("legends")
 
     def _is_blank(self, name: str) -> bool:
         return self.cmp.theme.T.is_blank(name)
 
     def _move_artists(self, spaces: CompositionSideSpaces):
         """
-        Move the annotations to their final positions
+        Move the annotations and legends to their final positions
         """
-        from ._plot_layout_items import _position_plot_labels
+        from ._plot_layout_items import (
+            _position_plot_labels,
+            set_legends_position,
+        )
 
-        _position_plot_labels(spaces.cmp.figure, self.cmp.theme, spaces, self)
+        # Only the root composition can have annotations (labels).
+        # So positioning them is a no-op. Skip the traversal entirely.
+        if self.is_root:
+            _position_plot_labels(
+                spaces.cmp.figure, self.cmp.theme, spaces, self
+            )
+        if self.legends:
+            set_legends_position(self.legends, spaces)

plotnine/_mpl/layout_manager/_composition_side_space.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from plotnine._mpl.layout_manager._layout_tree import LayoutTree
@@ -10,25 +11,66 @@
 
 if TYPE_CHECKING:
     from plotnine.composition._compose import Compose
+    from plotnine.iapi import outside_legend
 
 
 class _composition_side_space(_side_space):
     """
     Base class for the side space around a composition
+
+    The `plot_margin_*` figure-edge buffer is reserved only for the root
+    composition. Nested compositions sit inside that buffer and would
+    double-count it.
     """
 
     def __init__(self, items: CompositionLayoutItems):
         self.items = items
         self.gridspec = items.cmp._gridspec
         self._calculate()
 
+    @cached_property
+    def _legend_size(self) -> tuple[float, float]:
+        """
+        Return size of the side legend in figure coordinates
+        """
+        if not self.has_legend:
+            return (0, 0)
+
+        ol: outside_legend = getattr(self.items.legends, self.side)
+        return self.items.geometry.size(ol.box)
+
+    @cached_property
+    def legend_width(self) -> float:
+        return self._legend_size[0]
+
+    @cached_property
+    def legend_height(self) -> float:
+        return self._legend_size[1]
+
+    @property
+    def has_legend(self) -> bool:
+        """
+        Return True if this side has a legend to lay out
+        """
+        if not self.items.legends:
+            return False
+        return getattr(self.items.legends, self.side, None) is not None
+
 
 class composition_left_space(_composition_side_space):
     plot_margin: float = 0
+    legend: float = 0
+    legend_box_spacing: float = 0
 
     def _calculate(self):
-        theme = self.items.cmp.theme
-        self.plot_margin = theme.getp("plot_margin_left")
+        items = self.items
+        theme = items.cmp.theme
+        if items.is_root:
+            self.plot_margin = theme.getp("plot_margin_left")
+
+        if items.legends and items.legends.left:
+            self.legend = self.legend_width
+            self.legend_box_spacing = theme.getp("legend_box_spacing")
 
     @property
     def offset(self) -> float:
@@ -81,10 +123,18 @@ class composition_right_space(_composition_side_space):
     """
 
     plot_margin: float = 0
+    legend: float = 0
+    legend_box_spacing: float = 0
 
     def _calculate(self):
-        theme = self.items.cmp.theme
-        self.plot_margin = theme.getp("plot_margin_right")
+        items = self.items
+        theme = items.cmp.theme
+        if items.is_root:
+            self.plot_margin = theme.getp("plot_margin_right")
+
+        if items.legends and items.legends.right:
+            self.legend = self.legend_width
+            self.legend_box_spacing = theme.getp("legend_box_spacing")
 
     @property
     def offset(self):
@@ -143,6 +193,8 @@ class composition_top_space(_composition_side_space):
     plot_subtitle_margin_top: float = 0
     plot_subtitle: float = 0
     plot_subtitle_margin_bottom: float = 0
+    legend: float = 0
+    legend_box_spacing: float = 0
 
     def _calculate(self):
         items = self.items
@@ -151,7 +203,8 @@ def _calculate(self):
         W, H = theme.getp("figure_size")
         F = W / H
 
-        self.plot_margin = theme.getp("plot_margin_top") * F
+        if items.is_root:
+            self.plot_margin = theme.getp("plot_margin_top") * F
 
         if items.plot_title:
             m = theme.get_margin("plot_title").fig
@@ -165,6 +218,10 @@ def _calculate(self):
             self.plot_subtitle = geometry.height(items.plot_subtitle)
             self.plot_subtitle_margin_bottom = m.b
 
+        if items.legends and items.legends.top:
+            self.legend = self.legend_height
+            self.legend_box_spacing = theme.getp("legend_box_spacing") * F
+
     @property
     def offset(self) -> float:
         """
@@ -225,6 +282,8 @@ class composition_bottom_space(_composition_side_space):
     plot_caption_margin_bottom: float = 0
     plot_caption: float = 0
     plot_caption_margin_top: float = 0
+    legend: float = 0
+    legend_box_spacing: float = 0
 
     def _calculate(self):
         items = self.items
@@ -233,13 +292,18 @@ def _calculate(self):
         W, H = theme.getp("figure_size")
         F = W / H
 
-        self.plot_margin = theme.getp("plot_margin_bottom") * F
+        if items.is_root:
+            self.plot_margin = theme.getp("plot_margin_bottom") * F
         if items.plot_footer:
             m = theme.get_margin("plot_footer").fig
             self.plot_footer_margin_bottom = m.b
             self.plot_footer = geometry.height(items.plot_footer)
             self.plot_footer_margin_top = m.t
 
+        if items.legends and items.legends.bottom:
+            self.legend = self.legend_height
+            self.legend_box_spacing = theme.getp("legend_box_spacing") * F
+
         if items.plot_caption:
             m = theme.get_margin("plot_caption").fig
             self.plot_caption_margin_bottom = m.b
@@ -307,14 +371,17 @@ class CompositionSideSpaces:
     """
     Compute the spaces required to layout the composition
 
-    This is meant for the top-most composition
+    Built for the top-most composition and additionally for any
+    nested composition that collects guides (`layout.guides ==
+    "collect"`) — those need their own legend positioned within
+    the area their parent allocates.
     """
 
-    def __init__(self, cmp: Compose):
+    def __init__(self, cmp: Compose, *, is_root: bool = True):
         self.cmp = cmp
         self.gridspec = cmp._gridspec
         self.sub_gridspec = cmp._sub_gridspec
-        self.items = CompositionLayoutItems(cmp)
+        self.items = CompositionLayoutItems(cmp, is_root=is_root)
 
         self.l = composition_left_space(self.items)
         """All subspaces to the left of the panels"""
@@ -328,18 +395,42 @@ def __init__(self, cmp: Compose):
         self.b = composition_bottom_space(self.items)
         """All subspaces below the bottom of the panels"""
 
-        self._create_plot_sidespaces()
+        # The root creates PlotSideSpaces for every leaf in the
+        # tree, and side-spaces for every collecting nested cmp.
+        # Nested instances skip both to avoid double-creation.
+        self._nested_owners: list[Compose] = []
+        if is_root:
+            self._create_plot_sidespaces()
+            for sub in cmp._walk_guide_owners():
+                if sub is not cmp:
+                    sub._sidespaces = CompositionSideSpaces(sub, is_root=False)
+                    self._nested_owners.append(sub)
         self.tree = LayoutTree.create(cmp)
 
+    @property
+    def owner(self) -> Compose:
+        """
+        The composition these side-spaces are calculated against
+        """
+        return self.cmp
+
     def arrange(self):
         """
         Resize composition and place artists in final positions
         """
         # We first resize the compositions gridspec so that the tree
         # algorithms can work with the final position and total area.
         self.resize_gridspec()
+        # Collecting nested cmps shrink their own outer 1×1 to make
+        # room for their legend BEFORE alignment runs, so the tree's
+        # `align_panels` sees the actual panel area and lines panels
+        # up across nested boundaries.
+        for sub in self._nested_owners:
+            sub._sidespaces.resize_gridspec()
         self.tree.arrange_layout()
         self.items._move_artists(self)
+        for sub in self._nested_owners:
+            sub._sidespaces.items._move_artists(sub._sidespaces)
         self._arrange_plots()
 
     def _arrange_plots(self):

plotnine/_mpl/layout_manager/_plot_layout_items.py

@@ -6,6 +6,7 @@
 from matplotlib.text import Text
 
 from plotnine._mpl.patches import StripTextPatch
+from plotnine.composition._compose import Compose
 from plotnine.exceptions import PlotnineError
 
 from ..utils import (
@@ -574,13 +575,21 @@ def _position_plot_labels(
     return justify
 
 
-def set_legends_position(legends: legend_artists, spaces: PlotSideSpaces):
+def set_legends_position(
+    legends: legend_artists,
+    spaces: PlotSideSpaces | CompositionSideSpaces,
+):
     """
-    Place legend on the figure and justify is a required
+    Place legends on the figure, justifying each as required
+
+    Works for both plot-level and composition-level legends. Both
+    side-space hierarchies expose an `owner` property — a `ggplot`
+    or `Compose` — which provides the `_sub_gridspec` to anchor
+    against and the `figure` to transform onto.
     """
-    panels_gs = spaces.plot._sub_gridspec
+    panels_gs = spaces.owner._sub_gridspec
     params = panels_gs.get_subplot_params()
-    transFigure = spaces.plot.figure.transFigure
+    transFigure = spaces.owner.figure.transFigure
 
     def set_position(
         aob: FlexibleAnchoredOffsetbox,
@@ -645,8 +654,13 @@ def set_position(
         y = spaces.b.y1("legend")
         set_position(legends.bottom.box, (x, y), (0, 0))
 
-    # Inside legends are placed using the panels coordinate system
+    # Inside legends are placed using the panels coordinate system.
+    # For a `Compose` owner with a `guide_area` host, the guides are
+    # rendered in the guide_areas panel, so we need that gridspec
     if legends.inside:
+        if isinstance(spaces.owner, Compose) and spaces.owner._guide_area:
+            panels_gs = spaces.owner._guide_area._sub_gridspec
+
         transPanels = panels_gs.to_transform()
         for l in legends.inside:
             set_position(l.box, l.position, l.justification, transPanels)

plotnine/_mpl/layout_manager/_plot_side_space.py

@@ -766,6 +766,13 @@ def __init__(self, plot: ggplot):
 
         self.W, self.H = plot.theme.getp("figure_size")
 
+    @property
+    def owner(self) -> ggplot:
+        """
+        The plot these side-spaces are calculated against
+        """
+        return self.plot
+
     def arrange(self):
         """
         Resize plot and place artists in final positions around the panels

plotnine/_mpl/layout_manager/_side_space.py

@@ -69,7 +69,8 @@ class _side_space(ABC):
 
     A *_space class does the book keeping for all the artists that may
     fall on that side of the panels. The same name may appear in multiple
-    side classes (e.g. legend).
+    side classes (e.g. legend). The expected naming convention for the
+    subclasses is  `*(left|right|top|bottom)_space`.
 
     The amount of space for each artist is computed in figure coordinates.
     """
@@ -89,7 +90,7 @@ def side(self) -> Side:
         """
         Side of the panel(s) that this class applies to
         """
-        return cast("Side", self.__class__.__name__.split("_")[0])
+        return cast("Side", self.__class__.__name__.split("_")[-2])
 
     @cached_property
     def parts(self) -> list[str]:

plotnine/composition/__init__.py

@@ -1,5 +1,6 @@
 from ._beside import Beside
 from ._compose import Compose
+from ._guide_area import guide_area
 from ._inset_element import inset_element
 from ._plot_annotation import plot_annotation
 from ._plot_layout import plot_layout
@@ -12,6 +13,7 @@
     "Stack",
     "Beside",
     "Wrap",
+    "guide_area",
     "inset_element",
     "plot_annotation",
     "plot_layout",