diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index cd9b5d467..1e7e69a70 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -379,3 +379,34 @@ def __dir__(): # Prevent "import ultraplot.figure" from clobbering the top-level callable. install_module_proxy(sys.modules.get(__name__)) + + +def _patch_seaborn_move_legend(): + """ + 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") + 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 + + def _move_legend(obj, *args, **kwargs): + if isinstance(obj, SubplotGrid) and len(obj) == 1: + obj = obj[0] + 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 diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 19a429d15..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): @@ -3791,6 +3794,30 @@ 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) + + # 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): + 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)) + 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 = [ + pair for pair in remaining_pairs if pair[1] != label + ] + 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. @@ -3811,6 +3838,16 @@ def legend( **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): """ diff --git a/ultraplot/legend.py b/ultraplot/legend.py index 46c79fb7a..e262cb405 100644 --- a/ultraplot/legend.py +++ b/ultraplot/legend.py @@ -1729,6 +1729,44 @@ 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 + 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() + def _normalize_em_kwargs(kwargs: dict[str, Any], *, fontsize: float) -> dict[str, Any]: """ 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..2076028d0 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -287,6 +287,78 @@ 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_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_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.