Add public method set_loc() for Legend

cocolato · cocolato · commit 5aca651e3d34 · 2023-06-16T18:42:02.000Z
diff --git a/doc/users/next_whats_new/set_loc.rst b/doc/users/next_whats_new/set_loc.rst
@@ -0,0 +1,23 @@
+Add a public method to modify the location of ``Legend``
+--------------------------------------------------------
+
+`~matplotlib.legend.Legend` locations now can be tweaked after they've been defined.
+
+.. plot::
+    :include-source: true
+
+    from matplotlib import pyplot as plt
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+
+    x = list(range(-100, 101))
+    y = [i**2 for i in x]
+
+    ax.plot(x, y, label="f(x)")
+    ax.legend()
+    ax.get_legend().set_loc("right")
+    # Or
+    # ax.get_legend().set(loc="right")
+
+    plt.show()
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -330,6 +330,12 @@ def _update_bbox_to_anchor(self, loc_in_canvas):
     _legend_kw_doc_base)
 _docstring.interpd.update(_legend_kw_doc=_legend_kw_both_st)
 
+_legend_kw_set_loc_st = (
+    _loc_doc_base.format(parent='axes/figure',
+                         default=":rc:`legend.loc` for Axes, 'upper right' for Figure",
+                         best=_loc_doc_best, outside=_outside_doc))
+_docstring.interpd.update(_legend_kw_set_loc_doc=_legend_kw_set_loc_st)
+
 
 class Legend(Artist):
     """
@@ -503,58 +509,6 @@ def val_or_rc(val, rc_name):
             )
         self.parent = parent
 
-        loc0 = loc
-        self._loc_used_default = loc is None
-        if loc is None:
-            loc = mpl.rcParams["legend.loc"]
-            if not self.isaxes and loc in [0, 'best']:
-                loc = 'upper right'
-
-        type_err_message = ("loc must be string, coordinate tuple, or"
-                            f" an integer 0-10, not {loc!r}")
-
-        # handle outside legends:
-        self._outside_loc = None
-        if isinstance(loc, str):
-            if loc.split()[0] == 'outside':
-                # strip outside:
-                loc = loc.split('outside ')[1]
-                # strip "center" at the beginning
-                self._outside_loc = loc.replace('center ', '')
-                # strip first
-                self._outside_loc = self._outside_loc.split()[0]
-                locs = loc.split()
-                if len(locs) > 1 and locs[0] in ('right', 'left'):
-                    # locs doesn't accept "left upper", etc, so swap
-                    if locs[0] != 'center':
-                        locs = locs[::-1]
-                    loc = locs[0] + ' ' + locs[1]
-            # check that loc is in acceptable strings
-            loc = _api.check_getitem(self.codes, loc=loc)
-        elif np.iterable(loc):
-            # coerce iterable into tuple
-            loc = tuple(loc)
-            # validate the tuple represents Real coordinates
-            if len(loc) != 2 or not all(isinstance(e, numbers.Real) for e in loc):
-                raise ValueError(type_err_message)
-        elif isinstance(loc, int):
-            # validate the integer represents a string numeric value
-            if loc < 0 or loc > 10:
-                raise ValueError(type_err_message)
-        else:
-            # all other cases are invalid values of loc
-            raise ValueError(type_err_message)
-
-        if self.isaxes and self._outside_loc:
-            raise ValueError(
-                f"'outside' option for loc='{loc0}' keyword argument only "
-                "works for figure legends")
-
-        if not self.isaxes and loc == 0:
-            raise ValueError(
-                "Automatic legend placement (loc='best') not implemented for "
-                "figure legend")
-
         self._mode = mode
         self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform)
 
@@ -598,9 +552,8 @@ def val_or_rc(val, rc_name):
         # init with null renderer
         self._init_legend_box(handles, labels, markerfirst)
 
-        tmp = self._loc_used_default
-        self._set_loc(loc)
-        self._loc_used_default = tmp  # ignore changes done by _set_loc
+        # Set legend location
+        self.set_loc(loc)
 
         # figure out title font properties:
         if title_fontsize is not None and title_fontproperties is not None:
@@ -686,6 +639,73 @@ def _set_artist_props(self, a):
 
         a.set_transform(self.get_transform())
 
+    @_docstring.dedent_interpd
+    def set_loc(self, loc=None):
+        """
+        Set the location of the legend.
+
+        .. versionadded:: 3.8
+
+        Parameters
+        ----------
+        %(_legend_kw_set_loc_doc)s
+        """
+        loc0 = loc
+        self._loc_used_default = loc is None
+        if loc is None:
+            loc = mpl.rcParams["legend.loc"]
+            if not self.isaxes and loc in [0, 'best']:
+                loc = 'upper right'
+
+        type_err_message = ("loc must be string, coordinate tuple, or"
+                            f" an integer 0-10, not {loc!r}")
+
+        # handle outside legends:
+        self._outside_loc = None
+        if isinstance(loc, str):
+            if loc.split()[0] == 'outside':
+                # strip outside:
+                loc = loc.split('outside ')[1]
+                # strip "center" at the beginning
+                self._outside_loc = loc.replace('center ', '')
+                # strip first
+                self._outside_loc = self._outside_loc.split()[0]
+                locs = loc.split()
+                if len(locs) > 1 and locs[0] in ('right', 'left'):
+                    # locs doesn't accept "left upper", etc, so swap
+                    if locs[0] != 'center':
+                        locs = locs[::-1]
+                    loc = locs[0] + ' ' + locs[1]
+            # check that loc is in acceptable strings
+            loc = _api.check_getitem(self.codes, loc=loc)
+        elif np.iterable(loc):
+            # coerce iterable into tuple
+            loc = tuple(loc)
+            # validate the tuple represents Real coordinates
+            if len(loc) != 2 or not all(isinstance(e, numbers.Real) for e in loc):
+                raise ValueError(type_err_message)
+        elif isinstance(loc, int):
+            # validate the integer represents a string numeric value
+            if loc < 0 or loc > 10:
+                raise ValueError(type_err_message)
+        else:
+            # all other cases are invalid values of loc
+            raise ValueError(type_err_message)
+
+        if self.isaxes and self._outside_loc:
+            raise ValueError(
+                f"'outside' option for loc='{loc0}' keyword argument only "
+                "works for figure legends")
+
+        if not self.isaxes and loc == 0:
+            raise ValueError(
+                "Automatic legend placement (loc='best') not implemented for "
+                "figure legend")
+
+        tmp = self._loc_used_default
+        self._set_loc(loc)
+        self._loc_used_default = tmp  # ignore changes done by _set_loc
+
     def _set_loc(self, loc):
         # find_offset function will be provided to _legend_box and
         # _legend_box will draw itself at the location of the return
diff --git a/lib/matplotlib/legend.pyi b/lib/matplotlib/legend.pyi
@@ -118,6 +118,7 @@ class Legend(Artist):
     def get_texts(self) -> list[Text]: ...
     def set_alignment(self, alignment: Literal["center", "left", "right"]) -> None: ...
     def get_alignment(self) -> Literal["center", "left", "right"]: ...
+    def set_loc(self, loc: str | tuple[float, float] | int | None = ...) -> None: ...
     def set_title(
         self, title: str, prop: FontProperties | str | pathlib.Path | None = ...
     ) -> None: ...
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -755,6 +755,26 @@ def test_legend_alignment(alignment):
     assert leg.get_alignment() == alignment
 
 
+@pytest.mark.parametrize('loc', ('center', 'best',))
+def test_ax_legend_set_loc(loc):
+    fig, ax = plt.subplots()
+    ax.plot(range(10), label='test')
+    leg = ax.legend()
+    leg.set_loc(loc)
+    assert leg._get_loc() == mlegend.Legend.codes[loc]
+
+
+@pytest.mark.parametrize('loc', ('outside right', 'right',))
+def test_fig_legend_set_loc(loc):
+    fig, ax = plt.subplots()
+    ax.plot(range(10), label='test')
+    leg = fig.legend()
+    leg.set_loc(loc)
+
+    loc = loc.split()[1] if loc.startswith("outside") else loc
+    assert leg._get_loc() == mlegend.Legend.codes[loc]
+
+
 @pytest.mark.parametrize('alignment', ('center', 'left', 'right'))
 def test_legend_set_alignment(alignment):
     fig, ax = plt.subplots()
