Move facet-strip sizing and placement into the layout manager

plotnine/_mpl/layout_manager/_plot_layout_items.py

@@ -1,18 +1,19 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from itertools import chain
 from typing import TYPE_CHECKING
 
 from matplotlib.text import Text
 
-from plotnine._mpl.patches import StripTextPatch
-from plotnine._utils import side_artists
+from plotnine._utils import ha_as_float, side_artists, va_as_float
 from plotnine.composition._compose import Compose
 from plotnine.exceptions import PlotnineError
 
 from ..utils import (
     ArtistGeometry,
     TextJustifier,
+    bbox_in_axes_space,
     get_subplotspecs,
     rel_position,
     resize_footer_background,
@@ -32,7 +33,7 @@
     from matplotlib.figure import Figure
     from matplotlib.lines import Line2D
     from matplotlib.patches import Rectangle
-    from matplotlib.transforms import Transform
+    from matplotlib.transforms import Bbox, Transform
 
     from plotnine import ggplot
     from plotnine._mpl.offsetbox import FlexibleAnchoredOffsetbox
@@ -67,6 +68,31 @@
     )
 
 
+@dataclass
+class StripSizing:
+    """
+    Theme inputs that fix a strip background's size and offset
+    """
+
+    margin: Margin
+    """Strip text margin with the units in lines"""
+
+    strip_align: float
+    """How far the background is offset from the panel edge"""
+
+    bg_x: float
+    """Left of the strip background in transAxes"""
+
+    bg_y: float
+    """Bottom of the strip background in transAxes"""
+
+    bg_width: float
+    """Width of the strip background in transAxes (top strips)"""
+
+    bg_height: float
+    """Height of the strip background in transAxes (right strips)"""
+
+
 class PlotLayoutItems:
     """
     Objects required to compute the layout
@@ -193,57 +219,109 @@ def axis_ticks_y(self, ax: Axes) -> Iterator[Tick]:
 
         return chain(major, minor)
 
+    def _strip_sizing(self, position: StripPosition) -> StripSizing:
+        """
+        Theme inputs that fix one strip's background size and offset
+
+        The keys read depend on the side the strip sits on.
+        """
+        theme = self.plot.theme
+        if position == "top":
+            return StripSizing(
+                margin=theme.getp(("strip_text_x", "margin")).to("lines"),
+                strip_align=theme.getp("strip_align_x"),
+                bg_x=theme.getp(("strip_text_x", "x"), 0),
+                bg_y=1,
+                bg_width=theme.getp(("strip_background_x", "width"), 1),
+                bg_height=0,
+            )
+        else:
+            return StripSizing(
+                margin=theme.getp(("strip_text_y", "margin")).to("lines"),
+                strip_align=theme.getp("strip_align_y"),
+                bg_x=1,
+                bg_y=theme.getp(("strip_text_y", "y"), 0),
+                bg_width=0,
+                bg_height=theme.getp(("strip_background_y", "height"), 1),
+            )
+
+    def strip_patch_bbox(
+        self, strip_text: StripText, scale: float = 1
+    ) -> Bbox:
+        """
+        Figure-space bounding box of one strip's background patch
+
+        The breadth (height for top strips, width for right strips) is
+        scaled by `scale` so the layout manager can equalise strips in
+        the same group.
+        """
+        from matplotlib.transforms import Bbox
+
+        sizing = self._strip_sizing(strip_text.position)
+        m = sizing.margin
+        text_bbox = self.geometry.bbox(strip_text)
+        ax_bbox = self.geometry.bbox(strip_text.ax)
+        W, H = self.plot.figure.bbox.width, self.plot.figure.bbox.height
+        line_height = strip_text._line_height(self.geometry.renderer)
+
+        x0 = rel_position(sizing.bg_x, 0, ax_bbox.x0, ax_bbox.x1)
+        y0 = rel_position(sizing.bg_y, 0, ax_bbox.y0, ax_bbox.y1)
+
+        if strip_text.position == "top":
+            margins = (m.b + m.t) * line_height / H
+            width = ax_bbox.width * sizing.bg_width
+            height = (text_bbox.height + margins) * scale
+            y0 += height * sizing.strip_align
+        else:
+            margins = (m.l + m.r) * line_height / W
+            height = ax_bbox.height * sizing.bg_height
+            width = (text_bbox.width + margins) * scale
+            x0 += width * sizing.strip_align
+        return Bbox.from_bounds(x0, y0, width, height)
+
     def strip_text_x_extra_height(self, position: StripPosition) -> float:
         """
         Height taken up by the top strips that is outside the panels
         """
         if not self.strip_text_x:
             return 0
 
-        artists = [
-            st.patch if st.patch.get_visible() else st
-            for st in self.strip_text_x
-            if st.patch.position == position
-        ]
-
         heights = []
+        for st in self.strip_text_x:
+            if st.position != position:
+                continue
+            strip_align = self._strip_sizing(st.position).strip_align
+            if st.patch.get_visible():
+                # The patch bounds are not yet set, so derive its natural
+                # height from the sizing inputs.
+                h = self.strip_patch_bbox(st).height
+            else:
+                h = self.geometry.height(st)
+            heights.append(max(h + h * strip_align, 0))
 
-        for a in artists:
-            info = (
-                a.text.draw_info
-                if isinstance(a, StripTextPatch)
-                else a.draw_info
-            )
-            h = self.geometry.height(a)
-            heights.append(max(h + h * info.strip_align, 0))
-
-        return max(heights)
+        return max(heights) if heights else 0
 
     def strip_text_y_extra_width(self, position: StripPosition) -> float:
         """
-        Width taken up by the top strips that is outside the panels
+        Width taken up by the right strips that is outside the panels
         """
         if not self.strip_text_y:
             return 0
 
-        artists = [
-            st.patch if st.patch.get_visible() else st
-            for st in self.strip_text_y
-            if st.patch.position == position
-        ]
-
         widths = []
+        for st in self.strip_text_y:
+            if st.position != position:
+                continue
+            strip_align = self._strip_sizing(st.position).strip_align
+            if st.patch.get_visible():
+                # The patch bounds are not yet set, so derive its natural
+                # width from the sizing inputs.
+                w = self.strip_patch_bbox(st).width
+            else:
+                w = self.geometry.width(st)
+            widths.append(max(w + w * strip_align, 0))
 
-        for a in artists:
-            info = (
-                a.text.draw_info
-                if isinstance(a, StripTextPatch)
-                else a.draw_info
-            )
-            w = self.geometry.width(a)
-            widths.append(max(w + w * info.strip_align, 0))
-
-        return max(widths)
+        return max(widths) if widths else 0
 
     def axis_ticks_x_max_height_at(
         self, location: AxesLocation, side: str
@@ -448,8 +526,7 @@ def _move_artists(self, spaces: PlotSideSpaces):
 
         self._adjust_axis_text_x(justify)
         self._adjust_axis_text_y(justify)
-        self._strip_text_x_background_equal_heights()
-        self._strip_text_y_background_equal_widths()
+        self._place_strip_backgrounds()
 
     def _adjust_axis_text_x(self, justify: TextJustifier):
         """
@@ -553,37 +630,113 @@ def to_horizontal_axis_dimensions(value: float, ax: Axes) -> float:
                     )
                     justify.horizontally(text, ha, low, high, width=width)
 
-    def _strip_text_x_background_equal_heights(self):
+    def _strip_breadth_scales(
+        self, group: list[StripText], breadth: Literal["height", "width"]
+    ) -> list[float]:
         """
-        Make the strip_text_x_backgrounds have equal heights
+        Per-strip factor that equalises the breadth across a group
 
-        The smaller heights are expanded to match the largest height
+        Each strip's natural breadth is grown to match the largest in
+        the group, so the backgrounds share a common height (top strips)
+        or width (right strips).
         """
-        if not self.strip_text_x:
-            return
-
-        heights = [
-            self.geometry.bbox(t.patch).height for t in self.strip_text_x
-        ]
-        max_height = max(heights)
-        relative_heights = [max_height / h for h in heights]
-        for text, scale in zip(self.strip_text_x, relative_heights):
-            text.patch.expand = scale
+        natural = [getattr(self.strip_patch_bbox(st), breadth) for st in group]
+        largest = max(natural)
+        return [largest / b for b in natural]
 
-    def _strip_text_y_background_equal_widths(self):
+    def _place_strip_backgrounds(self):
         """
-        Make the strip_text_y_backgrounds have equal widths
-
-        The smaller widths are expanded to match the largest width
+        Fix each strip background at its final bounds and place its text
         """
-        if not self.strip_text_y:
-            return
+        groups: tuple[
+            tuple[list[StripText], Literal["height", "width"]], ...
+        ] = (
+            (self.strip_text_x or [], "height"),
+            (self.strip_text_y or [], "width"),
+        )
+        for group, breadth in groups:
+            if not group:
+                continue
+            scales = self._strip_breadth_scales(group, breadth)
+            for st, scale in zip(group, scales):
+                bbox = self.strip_patch_bbox(st, scale)
+                st.patch.set_bounds(bbox.bounds)
+                st.patch.set_transform(self.plot.figure.transFigure)
+                self._place_strip_text(st)
+
+    def _place_strip_text(self, st: StripText):
+        """
+        Justify the strip text within its final background bounds
+        """
+        theme = self.plot.theme
+        position = st.position
+        ax = st.ax
+        renderer = self.geometry.renderer
+        sizing = self._strip_sizing(position)
+        m = sizing.margin
+
+        patch_bbox = bbox_in_axes_space(st.patch, ax, renderer)
+        text_bbox = bbox_in_axes_space(st, ax, renderer)
+
+        if position == "top":
+            ha = theme.getp(("strip_text_x", "ha"), "center")
+            va = theme.getp(("strip_text_x", "va"), "center")
+            rel_x, rel_y = ha_as_float(ha), va_as_float(va)
+
+            # line_height and margins in axes space
+            line_height = st._line_height(renderer) / ax.bbox.height
+
+            x = (
+                # Justify horizontally within the strip_background
+                rel_position(
+                    rel_x,
+                    text_bbox.width + (line_height * (m.l + m.r)),
+                    patch_bbox.x0,
+                    patch_bbox.x1,
+                )
+                + (m.l * line_height)
+                + text_bbox.width / 2
+            )
+            # Setting the y position based on the bounding box is wrong
+            y = (
+                rel_position(
+                    rel_y,
+                    text_bbox.height,
+                    patch_bbox.y0 + m.b * line_height,
+                    patch_bbox.y1 - m.t * line_height,
+                )
+                + text_bbox.height / 2
+            )
+        else:  # "right"
+            ha = theme.getp(("strip_text_y", "ha"), "center")
+            va = theme.getp(("strip_text_y", "va"), "center")
+            rel_x, rel_y = ha_as_float(ha), va_as_float(va)
+
+            # line_height in axes space
+            line_height = st._line_height(renderer) / ax.bbox.width
+
+            x = (
+                rel_position(
+                    rel_x,
+                    text_bbox.width,
+                    patch_bbox.x0 + m.l * line_height,
+                    patch_bbox.x1 - m.r * line_height,
+                )
+                + text_bbox.width / 2
+            )
+            y = (
+                # Justify vertically within the strip_background
+                rel_position(
+                    rel_y,
+                    text_bbox.height + ((m.b + m.t) * line_height),
+                    patch_bbox.y0,
+                    patch_bbox.y1,
+                )
+                + (m.b * line_height)
+                + text_bbox.height / 2
+            )
 
-        widths = [self.geometry.bbox(t.patch).width for t in self.strip_text_y]
-        max_width = max(widths)
-        relative_widths = [max_width / w for w in widths]
-        for text, scale in zip(self.strip_text_y, relative_widths):
-            text.patch.expand = scale
+        st.set_position((x, y))
 
 
 def _text_is_visible(text: Text) -> bool: