From 9136d0d64723d6efcd95115d4305dbfc75ece0ab Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 13 Aug 2026 22:13:52 +1000 Subject: [PATCH 1/9] feat: add remove to legend --- ultraplot/legend.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/ultraplot/legend.py b/ultraplot/legend.py index 46c79fb7a..a6ed93f29 100644 --- a/ultraplot/legend.py +++ b/ultraplot/legend.py @@ -1729,6 +1729,25 @@ def set_loc(self, loc=None): where, type = old_loc self.axes._legend_dict[(loc, type)] = value + def remove(self): + """ + Remove the legend and sync Ultraplot guide tracking state. + + Matplotlib's base ``Legend.remove`` leaves Ultraplot's internal + ``_legend_dict`` and ``legend_`` pointers untouched. When callers + remove a legend (e.g., ``sns.move_legend``), stale entries can keep + showing old legends alongside newly added ones. Keep both systems in + sync before delegating to Matplotlib's removal logic. + """ + ax = self.axes + if ax is not None and getattr(ax, "_legend_dict", None) is not None: + for loc_align, value in tuple(ax._legend_dict.items()): + if value is self: + ax._legend_dict.pop(loc_align, None) + if ax is not None and getattr(ax, "legend_", None) is self: + ax.legend_ = None + return super().remove() + def _normalize_em_kwargs(kwargs: dict[str, Any], *, fontsize: float) -> dict[str, Any]: """ From 6fafc21450234d149ec1fb5bba5b416dd6521937 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 13 Aug 2026 22:24:43 +1000 Subject: [PATCH 2/9] Fix legend removal for seaborn move_legend compatibility --- ultraplot/legend.py | 19 ++++++++++++++++++ ultraplot/tests/test_integration.py | 31 +++++++++++++++++++++++++++++ ultraplot/tests/test_legend.py | 21 +++++++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/ultraplot/legend.py b/ultraplot/legend.py index a6ed93f29..e262cb405 100644 --- a/ultraplot/legend.py +++ b/ultraplot/legend.py @@ -1746,6 +1746,25 @@ def remove(self): ax._legend_dict.pop(loc_align, None) if ax is not None and getattr(ax, "legend_", None) is self: ax.legend_ = None + if getattr(self, "_remove_method", None) is None: + try: + self.set_visible(False) + except Exception: + pass + try: + if getattr(self, "_legend_box", None) is not None: + self._legend_box.set_visible(False) + except Exception: + pass + try: + for child in list(self.get_children()): + try: + child.set_visible(False) + except Exception: + pass + except Exception: + pass + return None return super().remove() diff --git a/ultraplot/tests/test_integration.py b/ultraplot/tests/test_integration.py index c82fd38d6..6f2589c06 100644 --- a/ultraplot/tests/test_integration.py +++ b/ultraplot/tests/test_integration.py @@ -72,6 +72,37 @@ def test_user_labeled_shading_appears_in_legend(): assert "CI band" in labels +def test_sns_move_legend_clears_ultraplot_legend_cache(): + """ + ``sns.move_legend`` should not leave stale entries in ``_legend_dict``. + """ + sns = pytest.importorskip("seaborn") + pd = pytest.importorskip("pandas") + + rng = np.random.default_rng(0) + fig, ax = uplt.subplots() + + df = pd.DataFrame(rng.normal(size=1000) + 3, columns=["Test Values"]) + with ax.external(): + sns.histplot(df, ax=ax, kde=True, legend=True) + + assert ax[0]._legend_dict + old_keys = set(ax[0]._legend_dict.keys()) + old_key = next(iter(old_keys)) + sns.move_legend(ax, "upper right") + + new_keys = set(ax[0]._legend_dict.keys()) + assert new_keys != set() + assert len(new_keys) == len(old_keys) + if old_key[0] != "upper right": + assert old_key not in new_keys + assert ("upper right", "center") in new_keys + + # there should be exactly one legend entry cached + assert len(ax[0]._legend_dict) == 1 + uplt.close(fig) + + @pytest.mark.mpl_image_compare def test_pint_quantities(rng): """ diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index a382d0e8b..2e8e3cd04 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -287,6 +287,27 @@ def test_sync_label_dict(rng): uplt.close(fig) +def test_legend_remove_clears_internal_dict_state(): + """ + Removing a legend should clear Ultraplot guide-tracking state. + + This prevents stale legends from staying registered when wrappers like + ``sns.move_legend`` remove and recreate legends. + """ + fig, ax = uplt.subplots() + ax.plot([0, 1, 2], label="line") + leg = ax.legend(loc="lower right") + + # Confirm the new legend is tracked. + assert any(v is leg for v in ax[0]._legend_dict.values()) + + # Remove it directly and verify Ultraplot state is cleaned up. + leg.remove() + assert not any(v is leg for v in ax[0]._legend_dict.values()) + assert ax[0].legend_ is None + uplt.close(fig) + + def test_external_mode_defers_on_the_fly_legend(): """ External mode should defer on-the-fly legend creation until explicitly requested. From 4045943b2921df16d868d77395f0764e8fbc73a4 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 13 Aug 2026 22:48:13 +1000 Subject: [PATCH 3/9] Test fallback remove path when _remove_method is absent --- ultraplot/tests/test_legend.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index 2e8e3cd04..0d10fb7b2 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -308,6 +308,25 @@ def test_legend_remove_clears_internal_dict_state(): uplt.close(fig) +def test_legend_remove_without_remove_method_uses_visibility_fallback(): + """ + If a legend does not expose a private remove method, ``remove`` should still + fall back to a safe hide-only path. + """ + fig, ax = uplt.subplots() + ax.plot([0, 1, 2], label="line") + leg = ax.legend(loc="lower right") + + # Force the compatibility path used by wrappers that do not implement a + # remove backend on the internal legend object. + setattr(leg, "_remove_method", None) + + assert leg.remove() is None + assert not any(v is leg for v in ax[0]._legend_dict.values()) + assert getattr(ax[0], "legend_", None) is None + uplt.close(fig) + + def test_external_mode_defers_on_the_fly_legend(): """ External mode should defer on-the-fly legend creation until explicitly requested. From c069ad9885281d6f4c8286f774452472dce63d7d Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:06:19 +1000 Subject: [PATCH 4/9] finalizing touches --- ultraplot/__init__.py | 35 ++++++++++++++++++++++++++++++ ultraplot/axes/base.py | 49 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index cd9b5d467..cf03f99b0 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -379,3 +379,38 @@ def __dir__(): # Prevent "import ultraplot.figure" from clobbering the top-level callable. install_module_proxy(sys.modules.get(__name__)) + + +def _patch_seaborn_move_legend(): + """Compatibility shim so seaborn can move legends on single-axis grids.""" + try: + import matplotlib.axes + import matplotlib.figure + import seaborn as sns + from seaborn.axisgrid import Grid + + from .gridspec import SubplotGrid + except Exception: + return + + move_legend = getattr(sns, "move_legend", None) + if not callable(move_legend) or getattr(move_legend, "_ultraplot", False): + return + + def _move_legend(obj, *args, **kwargs): + if isinstance(obj, SubplotGrid) and len(obj) == 1: + obj = obj[0] + if not isinstance( + obj, (matplotlib.axes.Axes, matplotlib.figure.Figure, Grid) + ): + # Keep semantics unchanged for unsupported objects. + return move_legend(obj, *args, **kwargs) + return move_legend(obj, *args, **kwargs) + + _move_legend.__name__ = "move_legend" + _move_legend.__doc__ = getattr(move_legend, "__doc__", None) + _move_legend._ultraplot = True + sns.move_legend = _move_legend + + +_patch_seaborn_move_legend() diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 19a429d15..eba1b980e 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3791,12 +3791,48 @@ def legend( loc = _translate_loc(loc, "legend", default=rc["legend.loc"]) align = kwargs.pop("align", None) align = _translate_loc(align, "align", default="center") + label_order = kwargs.pop("label_order", None) + order = label_order if isinstance(label_order, str) and label_order in {"C", "F"} else None + + # Seaborn passes ``label_order`` as a label sequence. UltraPlot's + # legend API does not expose this parameter, so we coerce it by + # reordering explicit handle/label inputs when possible. + if label_order is not None and not isinstance(label_order, str): + if labels is None and isinstance(handles, dict): + handles_dict = handles + handles, labels = zip( + *[ + (handles_dict[key], key) + for key in label_order + if key in handles_dict + ] + ) + handles, labels = list(handles), list(labels) + elif handles is not None and labels is not None: + label_to_handle = {label: handle for handle, label in zip(handles, labels)} + label_to_handle_keys = set(label_to_handle) + label_order_set = set(label_order) + ordered_pairs = [ + (label_to_handle[label], label) + for label in label_order + if label in label_to_handle_keys + ] + if ordered_pairs: + remaining_pairs = [ + (handle, label) + for handle, label in zip(handles, labels) + if label not in label_order_set + ] + handles, labels = zip(*((*ordered_pairs, *remaining_pairs),)) + handles, labels = list(handles), list(labels) # Either draw right now or queue up for later. Handles can be successively # added to a single location this way. Used for on-the-fly legends. queue = kwargs.pop("queue", False) if queue: - self._register_guide("legend", (handles, labels), (loc, align), **kwargs) + self._register_guide( + "legend", (handles, labels), (loc, align), order=order, **kwargs + ) else: return self._add_legend( handles, @@ -3808,9 +3844,20 @@ def legend( col=col, rows=rows, cols=cols, + order=order, **kwargs, ) + def add_legend(self, *args, **kwargs): + """ + Back-compatibility alias for older Matplotlib/Seaborn integrations that call + ``add_legend``. + + Newer code should call :meth:`legend`, but some callers still rely on this + Matplotlib-internal entry point. + """ + return self.legend(*args, **kwargs) + @docstring._snippet_manager def catlegend(self, categories, **kwargs): """ From 4b21870a14b5ff7e12447bc2efd13bf4ac60027f Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:10:27 +1000 Subject: [PATCH 5/9] Harden Seaborn legend compatibility --- ultraplot/__init__.py | 10 +------- ultraplot/axes/base.py | 45 +++++++++++++--------------------- ultraplot/tests/test_legend.py | 32 ++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 37 deletions(-) diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index cf03f99b0..bc23f18ab 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -384,13 +384,10 @@ def __dir__(): def _patch_seaborn_move_legend(): """Compatibility shim so seaborn can move legends on single-axis grids.""" try: - import matplotlib.axes - import matplotlib.figure import seaborn as sns - from seaborn.axisgrid import Grid from .gridspec import SubplotGrid - except Exception: + except ImportError: return move_legend = getattr(sns, "move_legend", None) @@ -400,11 +397,6 @@ def _patch_seaborn_move_legend(): def _move_legend(obj, *args, **kwargs): if isinstance(obj, SubplotGrid) and len(obj) == 1: obj = obj[0] - if not isinstance( - obj, (matplotlib.axes.Axes, matplotlib.figure.Figure, Grid) - ): - # Keep semantics unchanged for unsupported objects. - return move_legend(obj, *args, **kwargs) return move_legend(obj, *args, **kwargs) _move_legend.__name__ = "move_legend" diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index eba1b980e..d19e7d63b 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3792,47 +3792,37 @@ def legend( align = kwargs.pop("align", None) align = _translate_loc(align, "align", default="center") label_order = kwargs.pop("label_order", None) - order = label_order if isinstance(label_order, str) and label_order in {"C", "F"} else None # Seaborn passes ``label_order`` as a label sequence. UltraPlot's # legend API does not expose this parameter, so we coerce it by # reordering explicit handle/label inputs when possible. if label_order is not None and not isinstance(label_order, str): if labels is None and isinstance(handles, dict): - handles_dict = handles - handles, labels = zip( - *[ - (handles_dict[key], key) - for key in label_order - if key in handles_dict - ] - ) - handles, labels = list(handles), list(labels) - elif handles is not None and labels is not None: - label_to_handle = {label: handle for handle, label in zip(handles, labels)} - label_to_handle_keys = set(label_to_handle) - label_order_set = set(label_order) - ordered_pairs = [ - (label_to_handle[label], label) - for label in label_order - if label in label_to_handle_keys + pairs = [ + (handles[key], key) for key in label_order if key in handles ] - if ordered_pairs: + handles, labels = map(list, zip(*pairs)) if pairs else ([], []) + elif handles is not None and labels is not None: + pairs = list(zip(handles, labels)) + ordered_pairs = [] + remaining_pairs = list(pairs) + for label in label_order: + matches = [pair for pair in remaining_pairs if pair[1] == label] + ordered_pairs.extend(matches) remaining_pairs = [ - (handle, label) - for handle, label in zip(handles, labels) - if label not in label_order_set + pair for pair in remaining_pairs if pair[1] != label ] - handles, labels = zip(*((*ordered_pairs, *remaining_pairs),)) - handles, labels = list(handles), list(labels) + pairs = [*ordered_pairs, *remaining_pairs] + handles, labels = map(list, zip(*pairs)) if pairs else ([], []) + + if isinstance(label_order, str) and label_order in {"C", "F"}: + kwargs["order"] = label_order # Either draw right now or queue up for later. Handles can be successively # added to a single location this way. Used for on-the-fly legends. queue = kwargs.pop("queue", False) if queue: - self._register_guide( - "legend", (handles, labels), (loc, align), order=order, **kwargs - ) + self._register_guide("legend", (handles, labels), (loc, align), **kwargs) else: return self._add_legend( handles, @@ -3844,7 +3834,6 @@ def legend( col=col, rows=rows, cols=cols, - order=order, **kwargs, ) diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index 0d10fb7b2..e6db940ea 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -327,6 +327,38 @@ def test_legend_remove_without_remove_method_uses_visibility_fallback(): uplt.close(fig) +def test_legend_label_order_preserves_duplicate_labels(): + """A Seaborn-style label order should retain every matching handle.""" + fig, ax = uplt.subplots() + first, = ax.plot([0, 1], label="duplicate") + second, = ax.plot([1, 0], label="duplicate") + third, = ax.plot([0.5, 0.5], label="other") + + legend = ax.legend( + [first, second, third], + ["duplicate", "duplicate", "other"], + label_order=["other", "duplicate"], + ) + + assert [text.get_text() for text in legend.get_texts()] == [ + "other", + "duplicate", + "duplicate", + ] + uplt.close(fig) + + +def test_legend_label_order_allows_no_matching_dict_labels(): + """An empty Seaborn-style ordering should create an empty legend cleanly.""" + fig, ax = uplt.subplots() + line, = ax.plot([0, 1], label="line") + + legend = ax.legend({"line": line}, label_order=["missing"]) + + assert legend.get_texts() == [] + uplt.close(fig) + + def test_external_mode_defers_on_the_fly_legend(): """ External mode should defer on-the-fly legend creation until explicitly requested. From fd22ff248a9d6e4d645917888ba54b6ef5e98ba7 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:11:17 +1000 Subject: [PATCH 6/9] Black --- ultraplot/axes/base.py | 4 +--- ultraplot/tests/test_legend.py | 8 ++++---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index d19e7d63b..7ecd9486e 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3798,9 +3798,7 @@ def legend( # reordering explicit handle/label inputs when possible. if label_order is not None and not isinstance(label_order, str): if labels is None and isinstance(handles, dict): - pairs = [ - (handles[key], key) for key in label_order if key in handles - ] + pairs = [(handles[key], key) for key in label_order if key in handles] handles, labels = map(list, zip(*pairs)) if pairs else ([], []) elif handles is not None and labels is not None: pairs = list(zip(handles, labels)) diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index e6db940ea..2076028d0 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -330,9 +330,9 @@ def test_legend_remove_without_remove_method_uses_visibility_fallback(): def test_legend_label_order_preserves_duplicate_labels(): """A Seaborn-style label order should retain every matching handle.""" fig, ax = uplt.subplots() - first, = ax.plot([0, 1], label="duplicate") - second, = ax.plot([1, 0], label="duplicate") - third, = ax.plot([0.5, 0.5], label="other") + (first,) = ax.plot([0, 1], label="duplicate") + (second,) = ax.plot([1, 0], label="duplicate") + (third,) = ax.plot([0.5, 0.5], label="other") legend = ax.legend( [first, second, third], @@ -351,7 +351,7 @@ def test_legend_label_order_preserves_duplicate_labels(): def test_legend_label_order_allows_no_matching_dict_labels(): """An empty Seaborn-style ordering should create an empty legend cleanly.""" fig, ax = uplt.subplots() - line, = ax.plot([0, 1], label="line") + (line,) = ax.plot([0, 1], label="line") legend = ax.legend({"line": line}, label_order=["missing"]) From 06d91d7b8770dc20f7fccc2a03c7cf33759b81eb Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:14:07 +1000 Subject: [PATCH 7/9] Fix eager import and patch --- ultraplot/__init__.py | 11 +++++------ ultraplot/axes/base.py | 3 +++ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index bc23f18ab..543876f45 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -383,13 +383,14 @@ def __dir__(): def _patch_seaborn_move_legend(): """Compatibility shim so seaborn can move legends on single-axis grids.""" - try: - import seaborn as sns + import sys - from .gridspec import SubplotGrid - except ImportError: + sns = sys.modules.get("seaborn") + if sns is None: return + from .gridspec import SubplotGrid + move_legend = getattr(sns, "move_legend", None) if not callable(move_legend) or getattr(move_legend, "_ultraplot", False): return @@ -404,5 +405,3 @@ def _move_legend(obj, *args, **kwargs): _move_legend._ultraplot = True sns.move_legend = _move_legend - -_patch_seaborn_move_legend() diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 7ecd9486e..b6b9d0781 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -873,6 +873,9 @@ def external(self, value=True): """ Context manager toggling external mode during the block. """ + from .. import _patch_seaborn_move_legend + + _patch_seaborn_move_legend() return _ExternalModeMixin._ExternalContext(self, value) def _in_external_context(self): From d796898fe517630c4642bbb70bd41188df713e49 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:14:19 +1000 Subject: [PATCH 8/9] Black --- ultraplot/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index 543876f45..bd5c002ad 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -404,4 +404,3 @@ def _move_legend(obj, *args, **kwargs): _move_legend.__doc__ = getattr(move_legend, "__doc__", None) _move_legend._ultraplot = True sns.move_legend = _move_legend - From ba6525384740ca534dabc5dcc572f2a3d7e4bdd3 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 14 Aug 2026 00:15:30 +1000 Subject: [PATCH 9/9] Document monkey patch --- ultraplot/__init__.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index bd5c002ad..1e7e69a70 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -382,7 +382,13 @@ def __dir__(): def _patch_seaborn_move_legend(): - """Compatibility shim so seaborn can move legends on single-axis grids.""" + """ + Let ``sns.move_legend(ax, ...)`` accept singleton :class:`SubplotGrid` objects. + + Seaborn only accepts native Matplotlib axes, figures, and its own grids. The + wrapper unwraps a singleton grid to its underlying axes; callers can avoid + this compatibility patch by passing ``ax[0]`` directly. + """ import sys sns = sys.modules.get("seaborn")