Fix inset colorbar frame reflow for refaspect (#593)

Makes the solver for colorbars consistent with inset colorbars by ensuring that the drawing pipeline is correctly run for the inset.

Author: cvanelteren

ultraplot/axes/base.py

@@ -3149,19 +3149,31 @@ def draw(self, renderer=None, *args, **kwargs):
             self._colorbar_fill.update_ticks(manual_only=True)  # only if needed
         if self._inset_parent is not None and self._inset_zoom:
             self.indicate_inset_zoom()
+        needs_inset_reflow = bool(getattr(self, "_inset_colorbar_needs_reflow", False))
+        has_inset_frame = bool(
+            getattr(self, "_inset_colorbar_frame", None) is not None
+            and getattr(self, "_inset_colorbar_obj", None)
+        )
         super().draw(renderer, *args, **kwargs)
-        if getattr(self, "_inset_colorbar_obj", None) and getattr(
-            self, "_inset_colorbar_needs_reflow", False
-        ):
-            self._inset_colorbar_needs_reflow = False
+        if has_inset_frame:
+            if not needs_inset_reflow:
+                needs_inset_reflow = _inset_colorbar_frame_needs_reflow(
+                    self._inset_colorbar_obj,
+                    labelloc=getattr(self, "_inset_colorbar_labelloc", None),
+                    renderer=renderer,
+                )
+        if has_inset_frame and needs_inset_reflow:
             _reflow_inset_colorbar_frame(
                 self._inset_colorbar_obj,
                 labelloc=getattr(self, "_inset_colorbar_labelloc", None),
                 ticklen=getattr(
                     self, "_inset_colorbar_ticklen", units(rc["tick.len"], "pt")
                 ),
+                renderer=renderer,
             )
-            self.figure.canvas.draw_idle()
+            self._inset_colorbar_needs_reflow = False
+            # Re-draw synchronously so the current render pass sees reflowed bounds.
+            super().draw(renderer, *args, **kwargs)
 
     def get_tightbbox(self, renderer, *args, **kwargs):
         # Perform extra post-processing steps
@@ -4581,11 +4593,71 @@ def _apply_inset_colorbar_layout(
         frame.set_bounds(*bounds_frame)
 
 
+def _inset_colorbar_frame_needs_reflow(colorbar, *, labelloc: str, renderer) -> bool:
+    cax = colorbar.ax
+    layout = getattr(cax, "_inset_colorbar_layout", None)
+    frame = getattr(cax, "_inset_colorbar_frame", None)
+    if not layout or frame is None:
+        return False
+
+    orientation = layout["orientation"]
+    loc = layout["loc"]
+    ticklocation = layout["ticklocation"]
+    labelloc_layout = labelloc if isinstance(labelloc, str) else ticklocation
+    bboxes = []
+
+    longaxis = _get_colorbar_long_axis(colorbar)
+    try:
+        bbox = longaxis.get_tightbbox(renderer)
+    except Exception:
+        bbox = None
+    if bbox is not None:
+        bboxes.append(bbox)
+
+    label_axis = _get_axis_for(
+        labelloc_layout, loc, orientation=orientation, ax=colorbar
+    )
+    if label_axis.label.get_text():
+        try:
+            bboxes.append(label_axis.label.get_window_extent(renderer=renderer))
+        except Exception:
+            pass
+
+    for artist in (
+        getattr(colorbar, "outline", None),
+        getattr(colorbar, "solids", None),
+        getattr(colorbar, "dividers", None),
+    ):
+        if artist is None:
+            continue
+        try:
+            bboxes.append(artist.get_window_extent(renderer=renderer))
+        except Exception:
+            pass
+
+    if not bboxes:
+        return False
+
+    x0 = min(bbox.x0 for bbox in bboxes)
+    y0 = min(bbox.y0 for bbox in bboxes)
+    x1 = max(bbox.x1 for bbox in bboxes)
+    y1 = max(bbox.y1 for bbox in bboxes)
+    frame_bbox = frame.get_window_extent(renderer=renderer)
+    tol = 1.0
+    return (
+        frame_bbox.x0 > x0 + tol
+        or frame_bbox.y0 > y0 + tol
+        or frame_bbox.x1 < x1 - tol
+        or frame_bbox.y1 < y1 - tol
+    )
+
+
 def _reflow_inset_colorbar_frame(
     colorbar,
     *,
     labelloc: str,
     ticklen: float,
+    renderer=None,
 ):
     cax = colorbar.ax
     layout = getattr(cax, "_inset_colorbar_layout", None)
@@ -4623,7 +4695,7 @@ def _reflow_inset_colorbar_frame(
         cb_width = width
         cb_height = length
 
-    renderer = cax.figure._get_renderer()
+    renderer = renderer or cax.figure._get_renderer()
     if hasattr(colorbar, "update_ticks"):
         colorbar.update_ticks(manual_only=True)
     bboxes = []

ultraplot/colorbar.py

@@ -347,7 +347,9 @@ def add(
             cax._inset_colorbar_obj = obj
             cax._inset_colorbar_labelloc = labelloc
             cax._inset_colorbar_ticklen = ticklen
-            _register_inset_colorbar_reflow(ax.figure)
+            has_frame = getattr(cax, "_inset_colorbar_frame", None) is not None
+            if has_frame:
+                _register_inset_colorbar_reflow(ax.figure)
         kw_outline = {"edgecolor": color, "linewidth": linewidth}
         if obj.outline is not None:
             obj.outline.update(kw_outline)
@@ -958,6 +960,7 @@ def _reflow_inset_colorbar_frame(
     *,
     labelloc: Optional[str],
     ticklen: float,
+    renderer=None,
 ):
     cax = colorbar.ax
     layout = getattr(cax, "_inset_colorbar_layout", None)
@@ -995,7 +998,7 @@ def _reflow_inset_colorbar_frame(
         cb_width = width
         cb_height = length
 
-    renderer = cax.figure._get_renderer()
+    renderer = renderer or cax.figure._get_renderer()
     if hasattr(colorbar, "update_ticks"):
         colorbar.update_ticks(manual_only=True)
     bboxes = []

ultraplot/tests/test_colorbar.py

@@ -97,6 +97,37 @@ def test_inset_colorbar_frame_wraps_label(rng, orientation, labelloc):
     assert frame_bbox.y1 >= label_bbox.y1 - tol
 
 
+def test_inset_colorbar_frame_wraps_label_with_refaspect(rng):
+    """
+    Inset colorbar frame should include the label when figure sizing is refaspect-driven.
+    """
+    from ultraplot.axes.base import _get_axis_for
+
+    fig, ax = uplt.subplots(refaspect=2)
+    data = rng.random((20, 30))
+    m = ax.pcolormesh(data, vmin=0, vmax=1)
+    cb = ax.colorbar(m, loc="ul", label="title", frameon=True)
+    fig.canvas.draw()
+
+    frame = cb.ax._inset_colorbar_frame
+    assert frame is not None
+
+    renderer = fig.canvas.get_renderer()
+    frame_bbox = frame.get_window_extent(renderer)
+    layout = cb.ax._inset_colorbar_layout
+    labelloc = cb.ax._inset_colorbar_labelloc
+    labelloc_layout = labelloc if isinstance(labelloc, str) else layout["ticklocation"]
+    label_axis = _get_axis_for(
+        labelloc_layout, layout["loc"], orientation=layout["orientation"], ax=cb
+    )
+    label_bbox = label_axis.label.get_window_extent(renderer)
+    tol = 1.0
+    assert frame_bbox.x0 <= label_bbox.x0 + tol
+    assert frame_bbox.x1 >= label_bbox.x1 - tol
+    assert frame_bbox.y0 <= label_bbox.y0 + tol
+    assert frame_bbox.y1 >= label_bbox.y1 - tol
+
+
 from itertools import product