feat: add remove to legend

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

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):
  """

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]:
  """

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):
  """

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.