Refactor UltraLegend builder into dedicated module

ultraplot/axes/base.py

@@ -71,30 +71,6 @@
 # A-b-c label string
 ABC_STRING = "abcdefghijklmnopqrstuvwxyz"
 
-# Legend align options
-ALIGN_OPTS = {
-    None: {
-        "center": "center",
-        "left": "center left",
-        "right": "center right",
-        "top": "upper center",
-        "bottom": "lower center",
-    },
-    "left": {
-        "top": "upper right",
-        "center": "center right",
-        "bottom": "lower right",
-    },
-    "right": {
-        "top": "upper left",
-        "center": "center left",
-        "bottom": "lower left",
-    },
-    "top": {"left": "lower left", "center": "lower center", "right": "lower right"},
-    "bottom": {"left": "upper left", "center": "upper center", "right": "upper right"},
-}
-
-
 # Projection docstring
 _proj_docstring = """
 proj, projection : \\
@@ -1486,148 +1462,38 @@ def _add_legend(
         cols: Optional[Union[int, Tuple[int, int]]] = None,
         **kwargs,
     ):
-        """
-        The driver function for adding axes legends.
-        """
-        # Parse input argument units
-        ncol = _not_none(ncols=ncols, ncol=ncol)
-        order = _not_none(order, "C")
-        frameon = _not_none(frame=frame, frameon=frameon, default=rc["legend.frameon"])
-        fontsize = _not_none(fontsize, rc["legend.fontsize"])
-        titlefontsize = _not_none(
-            title_fontsize=kwargs.pop("title_fontsize", None),
-            titlefontsize=titlefontsize,
-            default=rc["legend.title_fontsize"],
-        )
-        fontsize = _fontsize_to_pt(fontsize)
-        titlefontsize = _fontsize_to_pt(titlefontsize)
-        if order not in ("F", "C"):
-            raise ValueError(
-                f"Invalid order {order!r}. Please choose from "
-                "'C' (row-major, default) or 'F' (column-major)."
-            )
-
-        # Convert relevant keys to em-widths
-        for setting in rcsetup.EM_KEYS:  # em-width keys
-            pair = setting.split("legend.", 1)
-            if len(pair) == 1:
-                continue
-            _, key = pair
-            value = kwargs.pop(key, None)
-            if isinstance(value, str):
-                value = units(value, "em", fontsize=fontsize)
-            if value is not None:
-                kwargs[key] = value
-
-        # Generate and prepare the legend axes
-        if loc in ("fill", "left", "right", "top", "bottom"):
-            lax = self._add_guide_panel(
-                loc,
-                align,
-                width=width,
-                space=space,
-                pad=pad,
-                span=span,
-                row=row,
-                col=col,
-                rows=rows,
-                cols=cols,
-            )
-            kwargs.setdefault("borderaxespad", 0)
-            if not frameon:
-                kwargs.setdefault("borderpad", 0)
-            try:
-                kwargs["loc"] = ALIGN_OPTS[lax._panel_side][align]
-            except KeyError:
-                raise ValueError(f"Invalid align={align!r} for legend loc={loc!r}.")
-        else:
-            lax = self
-            pad = kwargs.pop("borderaxespad", pad)
-            kwargs["loc"] = loc  # simply pass to legend
-            kwargs["borderaxespad"] = units(pad, "em", fontsize=fontsize)
-
-        # Handle and text properties that are applied after-the-fact
-        # NOTE: Set solid_capstyle to 'butt' so line does not extend past error bounds
-        # shading in legend entry. This change is not noticable in other situations.
-        kw_frame, kwargs = lax._parse_frame("legend", **kwargs)
-        kw_text = {}
-        if fontcolor is not None:
-            kw_text["color"] = fontcolor
-        if fontweight is not None:
-            kw_text["weight"] = fontweight
-        kw_title = {}
-        if titlefontcolor is not None:
-            kw_title["color"] = titlefontcolor
-        if titlefontweight is not None:
-            kw_title["weight"] = titlefontweight
-        kw_handle = _pop_props(kwargs, "line")
-        kw_handle.setdefault("solid_capstyle", "butt")
-        kw_handle.update(handle_kw or {})
-
-        # Parse the legend arguments using axes for auto-handle detection
-        # TODO: Update this when we no longer use "filled panels" for outer legends
-        pairs, multi = lax._parse_legend_handles(
+        return plegend.UltraLegend(self).add(
             handles,
             labels,
+            loc=loc,
+            align=align,
+            width=width,
+            pad=pad,
+            space=space,
+            frame=frame,
+            frameon=frameon,
             ncol=ncol,
-            order=order,
-            center=center,
+            ncols=ncols,
             alphabetize=alphabetize,
+            center=center,
+            order=order,
+            label=label,
+            title=title,
+            fontsize=fontsize,
+            fontweight=fontweight,
+            fontcolor=fontcolor,
+            titlefontsize=titlefontsize,
+            titlefontweight=titlefontweight,
+            titlefontcolor=titlefontcolor,
+            handle_kw=handle_kw,
             handler_map=handler_map,
+            span=span,
+            row=row,
+            col=col,
+            rows=rows,
+            cols=cols,
+            **kwargs,
         )
-        title = _not_none(label=label, title=title)
-        kwargs.update(
-            {
-                "title": title,
-                "frameon": frameon,
-                "fontsize": fontsize,
-                "handler_map": handler_map,
-                "title_fontsize": titlefontsize,
-            }
-        )
-
-        # Add the legend and update patch properties
-        # TODO: Add capacity for categorical labels in a single legend like seaborn
-        # rather than manual handle overrides with multiple legends.
-        if multi:
-            objs = lax._parse_legend_centered(pairs, kw_frame=kw_frame, **kwargs)
-        else:
-            kwargs.update({key: kw_frame.pop(key) for key in ("shadow", "fancybox")})
-            objs = [lax._parse_legend_aligned(pairs, ncol=ncol, order=order, **kwargs)]
-            objs[0].legendPatch.update(kw_frame)
-        for obj in objs:
-            if hasattr(lax, "legend_") and lax.legend_ is None:
-                lax.legend_ = obj  # make first legend accessible with get_legend()
-            else:
-                lax.add_artist(obj)
-
-        # Update legend patch and elements
-        # WARNING: legendHandles only contains the *first* artist per legend because
-        # HandlerBase.legend_artist() called in Legend._init_legend_box() only
-        # returns the first artist. Instead we try to iterate through offset boxes.
-        for obj in objs:
-            obj.set_clip_on(False)  # needed for tight bounding box calculations
-            box = getattr(obj, "_legend_handle_box", None)
-            for obj in guides._iter_children(box):
-                if isinstance(obj, mtext.Text):
-                    kw = kw_text
-                else:
-                    kw = {
-                        key: val
-                        for key, val in kw_handle.items()
-                        if hasattr(obj, "set_" + key)
-                    }  # noqa: E501
-                    if hasattr(obj, "set_sizes") and "markersize" in kw_handle:
-                        kw["sizes"] = np.atleast_1d(kw_handle["markersize"])
-                obj.update(kw)
-
-        # Register location and return
-        if isinstance(objs[0], mpatches.FancyBboxPatch):
-            objs = objs[1:]
-        obj = objs[0] if len(objs) == 1 else tuple(objs)
-        self._register_guide("legend", obj, (loc, align))  # possibly replace another
-
-        return obj
 
     def _apply_title_above(self):
         """