From f5a377c6999d1ee239ff72633faff61b302a34ee Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 30 Apr 2026 06:25:18 +1000 Subject: [PATCH 1/2] expose semantic legend to figure --- ultraplot/figure.py | 244 +++++++++++++++++++++++++++++++++ ultraplot/tests/test_legend.py | 71 ++++++++++ 2 files changed, 315 insertions(+) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 334c61a44..22a946711 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -31,6 +31,7 @@ from . import axes as paxes from . import constructor from . import gridspec as pgridspec +from . import legend as plegend from .config import rc, rc_matplotlib from .internals import ( _not_none, @@ -3052,6 +3053,249 @@ def _update_super_title(self, title, **kwargs): if title is not None: self._suptitle.set_text(title) + @staticmethod + def _iter_semantic_legend_axes(candidate): + """ + Yield axes objects from nested axis containers. + """ + if candidate is None or isinstance(candidate, str): + return + if isinstance(candidate, maxes.Axes): + yield candidate + return + if np.iterable(candidate): + for item in candidate: + yield from Figure._iter_semantic_legend_axes(item) + + def _semantic_legend_axes(self, ax=None, ref=None): + """ + Pick an axes instance for semantic legend handle generation. + """ + for candidate in (ax, ref, self.axes): + for axis in self._iter_semantic_legend_axes(candidate): + return axis + raise RuntimeError( + "Figure semantic legend helpers require an existing axes. " + "Create an axes first or pass ax=... or ref=...." + ) + + def entrylegend( + self, + entries, + *, + line=None, + marker=None, + color=None, + linestyle=None, + linewidth=None, + markersize=None, + alpha=None, + markeredgecolor=None, + markeredgewidth=None, + markerfacecolor=None, + handle_kw=None, + add=True, + **legend_kwargs, + ): + """ + Build generic semantic legend entries and optionally add a figure legend. + """ + axes = self._semantic_legend_axes( + ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") + ) + handles, labels = plegend.UltraLegend(axes).entrylegend( + entries, + line=line, + marker=marker, + color=color, + linestyle=linestyle, + linewidth=linewidth, + markersize=markersize, + alpha=alpha, + markeredgecolor=markeredgecolor, + markeredgewidth=markeredgewidth, + markerfacecolor=markerfacecolor, + handle_kw=handle_kw, + add=False, + ) + if not add: + return handles, labels + return self.legend(handles, labels, **legend_kwargs) + + def catlegend( + self, + categories, + *, + colors=None, + markers=None, + line=None, + linestyle=None, + linewidth=None, + markersize=None, + alpha=None, + markeredgecolor=None, + markeredgewidth=None, + markerfacecolor=None, + handle_kw=None, + add=True, + **legend_kwargs, + ): + """ + Build categorical legend entries and optionally add a figure legend. + """ + axes = self._semantic_legend_axes( + ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") + ) + handles, labels = plegend.UltraLegend(axes).catlegend( + categories, + colors=colors, + markers=markers, + line=line, + linestyle=linestyle, + linewidth=linewidth, + markersize=markersize, + alpha=alpha, + markeredgecolor=markeredgecolor, + markeredgewidth=markeredgewidth, + markerfacecolor=markerfacecolor, + handle_kw=handle_kw, + add=False, + ) + if not add: + return handles, labels + return self.legend(handles, labels, **legend_kwargs) + + def sizelegend( + self, + levels, + *, + labels=None, + color=None, + marker=None, + area=None, + scale=None, + minsize=None, + fmt=None, + alpha=None, + markeredgecolor=None, + markeredgewidth=None, + markerfacecolor=None, + handle_kw=None, + add=True, + **legend_kwargs, + ): + """ + Build size legend entries and optionally add a figure legend. + """ + axes = self._semantic_legend_axes( + ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") + ) + handles, labels = plegend.UltraLegend(axes).sizelegend( + levels, + labels=labels, + color=color, + marker=marker, + area=area, + scale=scale, + minsize=minsize, + fmt=fmt, + alpha=alpha, + markeredgecolor=markeredgecolor, + markeredgewidth=markeredgewidth, + markerfacecolor=markerfacecolor, + handle_kw=handle_kw, + add=False, + ) + if not add: + return handles, labels + return self.legend(handles, labels, **legend_kwargs) + + def numlegend( + self, + levels=None, + *, + vmin=None, + vmax=None, + n=None, + cmap=None, + norm=None, + fmt=None, + facecolor=None, + edgecolor=None, + linewidth=None, + linestyle=None, + alpha=None, + handle_kw=None, + add=True, + **legend_kwargs, + ): + """ + Build numeric-color legend entries and optionally add a figure legend. + """ + axes = self._semantic_legend_axes( + ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") + ) + handles, labels = plegend.UltraLegend(axes).numlegend( + levels=levels, + vmin=vmin, + vmax=vmax, + n=n, + cmap=cmap, + norm=norm, + fmt=fmt, + facecolor=facecolor, + edgecolor=edgecolor, + linewidth=linewidth, + linestyle=linestyle, + alpha=alpha, + handle_kw=handle_kw, + add=False, + ) + if not add: + return handles, labels + return self.legend(handles, labels, **legend_kwargs) + + def geolegend( + self, + entries, + labels=None, + *, + country_reso=None, + country_territories=None, + country_proj=None, + handlesize=None, + facecolor=None, + edgecolor=None, + linewidth=None, + alpha=None, + fill=None, + add=True, + **legend_kwargs, + ): + """ + Build geometry legend entries and optionally add a figure legend. + """ + axes = self._semantic_legend_axes( + ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") + ) + handles, labels = plegend.UltraLegend(axes).geolegend( + entries, + labels=labels, + country_reso=country_reso, + country_territories=country_territories, + country_proj=country_proj, + handlesize=handlesize, + facecolor=facecolor, + edgecolor=edgecolor, + linewidth=linewidth, + alpha=alpha, + fill=fill, + add=False, + ) + if not add: + return handles, labels + return self.legend(handles, labels, **legend_kwargs) + @_clear_border_cache @docstring._concatenate_inherited @docstring._snippet_manager diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index d9971fe16..ab7cc2fba 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -660,6 +660,77 @@ def test_semantic_legend_rejects_labels_kwarg(builder, args, kwargs): uplt.close(fig) +@pytest.mark.parametrize( + "builder, args, kwargs", + ( + ( + "entrylegend", + ([{"label": "Trend", "line": True}, {"label": "Samples", "line": False}],), + {}, + ), + ("catlegend", (["A", "B"],), {"colors": ["red7", "blue7"]}), + ( + "sizelegend", + ([10, 50],), + {"labels": ["small", "large"], "color": "gray6"}, + ), + ("numlegend", tuple(), {"levels": [0, 1], "cmap": "viridis"}), + ( + "geolegend", + ([("Triangle", "triangle"), ("Hex", "hexagon")],), + {}, + ), + ), +) +def test_figure_semantic_legend_helpers(builder, args, kwargs): + fig, axs = uplt.subplots(ncols=2) + ax = axs[0] + figure_method = getattr(fig, builder) + axes_method = getattr(ax, builder) + + expected_handles, expected_labels = axes_method(*args, add=False, **kwargs) + leg = figure_method(*args, ref=axs, loc="bottom", title=builder, **kwargs) + + assert leg is not None + assert [text.get_text() for text in leg.get_texts()] == expected_labels + assert leg.get_title().get_text() == builder + assert len(leg.legend_handles) == len(expected_handles) + uplt.close(fig) + + +@pytest.mark.parametrize( + "builder, args, kwargs", + ( + ("entrylegend", ([{"label": "Trend", "line": True}],), {}), + ("catlegend", (["A", "B"],), {}), + ("sizelegend", ([10, 50],), {"labels": ["small", "large"]}), + ("numlegend", tuple(), {"levels": [0, 1]}), + ("geolegend", (["triangle"], ["Triangle"]), {}), + ), +) +def test_figure_semantic_legend_add_false_matches_axes(builder, args, kwargs): + fig, ax = uplt.subplots() + figure_method = getattr(fig, builder) + axes_method = getattr(ax, builder) + + fig_handles, fig_labels = figure_method(*args, add=False, **kwargs) + ax_handles, ax_labels = axes_method(*args, add=False, **kwargs) + + assert fig_labels == ax_labels + assert len(fig_handles) == len(ax_handles) + assert [handle.get_label() for handle in fig_handles] == [ + handle.get_label() for handle in ax_handles + ] + uplt.close(fig) + + +def test_figure_semantic_legend_without_axes_raises(): + fig = uplt.figure() + with pytest.raises(RuntimeError, match="require an existing axes"): + fig.catlegend(["A"], loc="right") + uplt.close(fig) + + def test_geo_legend_handlesize_scales_handle_box(): fig, ax = uplt.subplots() leg = ax.geolegend([("shape", "triangle")], loc="best", handlesize=2.0) From 02d1c5d104186d64570e7a8d46eb338e0a766494 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 30 Apr 2026 06:35:26 +1000 Subject: [PATCH 2/2] Add documentation --- docs/colorbars_legends.py | 58 +++++++- .../legends_colorbars/03_semantic_legends.py | 30 +++- ultraplot/figure.py | 133 +++++++++++++++++- 3 files changed, 212 insertions(+), 9 deletions(-) diff --git a/docs/colorbars_legends.py b/docs/colorbars_legends.py index 4044a9200..5b0582622 100644 --- a/docs/colorbars_legends.py +++ b/docs/colorbars_legends.py @@ -477,13 +477,19 @@ # # Legends usually annotate artists already drawn on an axes, but sometimes you need # standalone semantic keys (categories, size scales, color levels, or geometry types). -# UltraPlot provides helper methods that build these entries directly: +# UltraPlot provides helper methods that build these entries directly on both +# axes and figures: # # * :meth:`~ultraplot.axes.Axes.entrylegend` # * :meth:`~ultraplot.axes.Axes.catlegend` # * :meth:`~ultraplot.axes.Axes.sizelegend` # * :meth:`~ultraplot.axes.Axes.numlegend` # * :meth:`~ultraplot.axes.Axes.geolegend` +# * :meth:`~ultraplot.figure.Figure.entrylegend` +# * :meth:`~ultraplot.figure.Figure.catlegend` +# * :meth:`~ultraplot.figure.Figure.sizelegend` +# * :meth:`~ultraplot.figure.Figure.numlegend` +# * :meth:`~ultraplot.figure.Figure.geolegend` # # These helpers are useful whenever the legend should describe an encoding rather than # mirror artists that already happen to be drawn. In practice there are two distinct @@ -513,7 +519,8 @@ # # The helpers are intentionally composable. Each one accepts ``add=False`` and returns # ``(handles, labels)`` so you can merge semantic sections and pass the result through -# :meth:`~ultraplot.axes.Axes.legend` yourself. +# :meth:`~ultraplot.axes.Axes.legend` or :meth:`~ultraplot.figure.Figure.legend` +# yourself. # # .. code-block:: python # @@ -568,6 +575,27 @@ # # .. code-block:: python # +# # Add semantic legends around an entire subplot group. +# fig, axs = uplt.subplots(ncols=2) +# fig.catlegend( +# ["Control", "Treatment"], +# colors={"Control": "blue7", "Treatment": "red7"}, +# markers={"Control": "o", "Treatment": "^"}, +# ref=axs, +# loc="b", +# title="Group", +# ) +# fig.sizelegend( +# [10, 50, 200], +# labels=["Small", "Medium", "Large"], +# color="gray6", +# ref=axs, +# loc="r", +# title="Population", +# ) +# +# .. code-block:: python +# # # Compose multiple semantic helpers into one legend. # size_handles, size_labels = ax.sizelegend( # [10, 50, 200], @@ -685,6 +713,32 @@ ax.axis("off") +# %% +fig, axs = uplt.subplots(ncols=2, refwidth=2.8, share=False) +axs[0].scatter([0, 1, 2], [3, 1, 2], c=[0.2, 0.5, 0.8], s=[40, 120, 260]) +axs[1].scatter([0, 1, 2], [2, 3, 1], c=[0.8, 0.4, 0.1], s=[60, 90, 220]) +axs.format(title="Figure semantic legend helpers", grid=False) + +fig.catlegend( + ["Control", "Treatment"], + colors={"Control": "blue7", "Treatment": "red7"}, + markers={"Control": "o", "Treatment": "^"}, + ref=axs, + loc="bottom", + title="Group", + frameon=False, +) +fig.sizelegend( + [40, 120, 260], + labels=["Small", "Medium", "Large"], + color="gray6", + ref=axs, + loc="right", + title="Size scale", + frameon=False, +) + + # %% [raw] raw_mimetype="text/restructuredtext" # .. _ug_guides_decouple: # diff --git a/docs/examples/legends_colorbars/03_semantic_legends.py b/docs/examples/legends_colorbars/03_semantic_legends.py index a869b826e..ea58d5bb9 100644 --- a/docs/examples/legends_colorbars/03_semantic_legends.py +++ b/docs/examples/legends_colorbars/03_semantic_legends.py @@ -6,12 +6,12 @@ Why UltraPlot here? ------------------- -UltraPlot adds semantic legend helpers directly on axes: +UltraPlot adds semantic legend helpers on both axes and figures: ``entrylegend``, ``catlegend``, ``sizelegend``, ``numlegend``, and ``geolegend``. These are useful when you want legend meaning decoupled from plotted handles, or when you want a standalone semantic key that describes an encoding directly. -Key functions: :py:meth:`ultraplot.axes.Axes.entrylegend`, :py:meth:`ultraplot.axes.Axes.catlegend`, :py:meth:`ultraplot.axes.Axes.sizelegend`, :py:meth:`ultraplot.axes.Axes.numlegend`, :py:meth:`ultraplot.axes.Axes.geolegend`. +Key functions: :py:meth:`ultraplot.axes.Axes.entrylegend`, :py:meth:`ultraplot.axes.Axes.catlegend`, :py:meth:`ultraplot.axes.Axes.sizelegend`, :py:meth:`ultraplot.axes.Axes.numlegend`, :py:meth:`ultraplot.axes.Axes.geolegend`, :py:meth:`ultraplot.figure.Figure.entrylegend`, :py:meth:`ultraplot.figure.Figure.catlegend`, :py:meth:`ultraplot.figure.Figure.sizelegend`, :py:meth:`ultraplot.figure.Figure.numlegend`, :py:meth:`ultraplot.figure.Figure.geolegend`. See also -------- @@ -106,3 +106,29 @@ ) ax.axis("off") fig.show() + +# %% +fig, axs = uplt.subplots(ncols=2, refwidth=2.8, share=False) +axs[0].scatter([0, 1, 2], [3, 1, 2], c=[0.2, 0.5, 0.8], s=[40, 120, 260]) +axs[1].scatter([0, 1, 2], [2, 3, 1], c=[0.8, 0.4, 0.1], s=[60, 90, 220]) +axs.format(title="Figure semantic legend helpers") + +fig.catlegend( + ["Control", "Treatment"], + colors={"Control": "blue7", "Treatment": "red7"}, + markers={"Control": "o", "Treatment": "^"}, + ref=axs, + loc="bottom", + title="Group", + frameon=False, +) +fig.sizelegend( + [40, 120, 260], + labels=["Small", "Medium", "Large"], + color="gray6", + ref=axs, + loc="right", + title="Size scale", + frameon=False, +) +fig.show() diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 22a946711..1ca87384e 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -412,6 +412,124 @@ ) # noqa: E501 +# Figure semantic legend helpers +_figure_semantic_legend_common_docstring = """ +**legend_kwargs + Placement and legend styling keywords forwarded to + `~ultraplot.figure.Figure.legend` when ``add=True``. This includes figure legend + placement keywords like ``loc=``, ``ref=``, ``ax=``, ``rows=``, ``cols=``, and + ``span=``. Pass ``add=False`` to return ``(handles, labels)`` without drawing. +""" +docstring._snippet_manager["figure.semantic_legend_common"] = ( + _figure_semantic_legend_common_docstring +) + +_figure_entrylegend_docstring = """ +Build generic semantic legend entries and optionally add a figure legend. + +Parameters +---------- +entries + Entry specifications as handles, style dictionaries, or ``(label, spec)`` + pairs. + +Other parameters +---------------- +%(figure.semantic_legend_common)s + +Notes +----- +Handle generation currently reuses the semantic legend builder used by +`~ultraplot.axes.Axes.entrylegend`, then routes the final draw step through +`~ultraplot.figure.Figure.legend`. +""" +docstring._snippet_manager["figure.entrylegend"] = _figure_entrylegend_docstring + +_figure_catlegend_docstring = """ +Build categorical legend entries and optionally add a figure legend. + +Parameters +---------- +categories + Category labels used to generate legend handles. + +Other parameters +---------------- +%(figure.semantic_legend_common)s + +Notes +----- +Handle generation currently reuses the semantic legend builder used by +`~ultraplot.axes.Axes.catlegend`, then routes the final draw step through +`~ultraplot.figure.Figure.legend`. +""" +docstring._snippet_manager["figure.catlegend"] = _figure_catlegend_docstring + +_figure_sizelegend_docstring = """ +Build size legend entries and optionally add a figure legend. + +Parameters +---------- +levels + Numeric levels used to generate marker-size entries. + +Other parameters +---------------- +%(figure.semantic_legend_common)s + +Notes +----- +Handle generation currently reuses the semantic legend builder used by +`~ultraplot.axes.Axes.sizelegend`, then routes the final draw step through +`~ultraplot.figure.Figure.legend`. + +Pass ``labels=[...]`` or ``labels={level: label}`` to override the generated labels. +""" +docstring._snippet_manager["figure.sizelegend"] = _figure_sizelegend_docstring + +_figure_numlegend_docstring = """ +Build numeric-color legend entries and optionally add a figure legend. + +Parameters +---------- +levels + Numeric levels or number of levels. + +Other parameters +---------------- +%(figure.semantic_legend_common)s + +Notes +----- +Handle generation currently reuses the semantic legend builder used by +`~ultraplot.axes.Axes.numlegend`, then routes the final draw step through +`~ultraplot.figure.Figure.legend`. +""" +docstring._snippet_manager["figure.numlegend"] = _figure_numlegend_docstring + +_figure_geolegend_docstring = """ +Build geometry legend entries and optionally add a figure legend. + +Parameters +---------- +entries + Geometry entries (mapping, ``(label, geometry)`` pairs, or geometries). +labels + Optional labels for geometry sequences. + +Other parameters +---------------- +%(figure.semantic_legend_common)s + +Notes +----- +Handle generation currently reuses the semantic legend builder used by +`~ultraplot.axes.Axes.geolegend`, then routes the final draw step through +`~ultraplot.figure.Figure.legend`. +""" +docstring._snippet_manager["figure.geolegend"] = _figure_geolegend_docstring + + # Save docstring _save_docstring = """ Save the figure. @@ -3079,6 +3197,7 @@ def _semantic_legend_axes(self, ax=None, ref=None): "Create an axes first or pass ax=... or ref=...." ) + @docstring._snippet_manager def entrylegend( self, entries, @@ -3098,7 +3217,7 @@ def entrylegend( **legend_kwargs, ): """ - Build generic semantic legend entries and optionally add a figure legend. + %(figure.entrylegend)s """ axes = self._semantic_legend_axes( ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") @@ -3122,6 +3241,7 @@ def entrylegend( return handles, labels return self.legend(handles, labels, **legend_kwargs) + @docstring._snippet_manager def catlegend( self, categories, @@ -3141,7 +3261,7 @@ def catlegend( **legend_kwargs, ): """ - Build categorical legend entries and optionally add a figure legend. + %(figure.catlegend)s """ axes = self._semantic_legend_axes( ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") @@ -3165,6 +3285,7 @@ def catlegend( return handles, labels return self.legend(handles, labels, **legend_kwargs) + @docstring._snippet_manager def sizelegend( self, levels, @@ -3185,7 +3306,7 @@ def sizelegend( **legend_kwargs, ): """ - Build size legend entries and optionally add a figure legend. + %(figure.sizelegend)s """ axes = self._semantic_legend_axes( ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") @@ -3210,6 +3331,7 @@ def sizelegend( return handles, labels return self.legend(handles, labels, **legend_kwargs) + @docstring._snippet_manager def numlegend( self, levels=None, @@ -3230,7 +3352,7 @@ def numlegend( **legend_kwargs, ): """ - Build numeric-color legend entries and optionally add a figure legend. + %(figure.numlegend)s """ axes = self._semantic_legend_axes( ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref") @@ -3255,6 +3377,7 @@ def numlegend( return handles, labels return self.legend(handles, labels, **legend_kwargs) + @docstring._snippet_manager def geolegend( self, entries, @@ -3273,7 +3396,7 @@ def geolegend( **legend_kwargs, ): """ - Build geometry legend entries and optionally add a figure legend. + %(figure.geolegend)s """ axes = self._semantic_legend_axes( ax=legend_kwargs.get("ax"), ref=legend_kwargs.get("ref")