format the set_loc() api doc

cocolato · cocolato · commit 7132ab05a203 · 2023-06-15T10:01:39.000Z
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):
     """
@@ -634,19 +640,14 @@ def _set_artist_props(self, a):
 
         a.set_transform(self.get_transform())
 
-    def set_loc(self, loc):
+    @_docstring.dedent_interpd
+    def set_loc(self, loc=None):
         """
         Set the location of the legend.
 
         Parameters
         ----------
-        loc : str
-            {'upper left', 'upper right', 'lower left', 'lower right',
-            'upper center', 'lower center', 'center left', 'center right'
-            'center', 'right', 'best'}.
-            If a figure is using the constrained layout manager,
-            the string codes of the loc keyword argument can get better
-            layout behaviour using the prefix 'outside'.
+        %(_legend_kw_set_loc_doc)s
         """
         loc0 = loc
         self._loc_used_default = loc is None
diff --git a/lib/matplotlib/legend.pyi b/lib/matplotlib/legend.pyi
@@ -118,7 +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) -> None: ...
+    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,15 +755,26 @@ def test_legend_alignment(alignment):
     assert leg.get_alignment() == alignment
 
 
-@pytest.mark.parametrize('loc', ('center', 'best'))
-def test_legend_set_loc(loc):
+@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()
