ENH: Type the possible str legend locs as Literals

Instead of accepting any str, we only accept as set of predefined
literals. For simplicity, we don't distinguish the between allowed
positions for Axes legend and figure legend. It's still better to limit
the allowed range to the union of both rather than to accept abitrary
strings.

Author: timhoffm

lib/matplotlib/axes/_axes.pyi

@@ -36,7 +36,7 @@ from collections.abc import Callable, Iterable, Sequence
 from typing import Any, Literal, overload
 import numpy as np
 from numpy.typing import ArrayLike
-from matplotlib.typing import ColorType, MarkerType, LineStyleType
+from matplotlib.typing import ColorType, MarkerType, LegendLocType, LineStyleType
 import pandas as pd
 
 
@@ -65,13 +65,16 @@ class Axes(_AxesBase):
     @overload
     def legend(self) -> Legend: ...
     @overload
-    def legend(self, handles: Iterable[Artist | tuple[Artist, ...]], labels: Iterable[str], **kwargs) -> Legend: ...
+    def legend(self, handles: Iterable[Artist | tuple[Artist, ...]], labels: Iterable[str],
+               *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, *, handles: Iterable[Artist | tuple[Artist, ...]], **kwargs) -> Legend: ...
+    def legend(self, *, handles: Iterable[Artist | tuple[Artist, ...]],
+               loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, labels: Iterable[str], **kwargs) -> Legend: ...
+    def legend(self, labels: Iterable[str],
+               *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, **kwargs) -> Legend: ...
+    def legend(self, *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
 
     def inset_axes(
         self,

lib/matplotlib/figure.pyi

@@ -27,7 +27,7 @@ from matplotlib.text import Text
 from matplotlib.transforms import Affine2D, Bbox, BboxBase, Transform
 from mpl_toolkits.mplot3d import Axes3D
 
-from .typing import ColorType, HashableList
+from .typing import ColorType, HashableList, LegendLocType
 
 _T = TypeVar("_T")
 
@@ -152,13 +152,16 @@ class FigureBase(Artist):
     @overload
     def legend(self) -> Legend: ...
     @overload
-    def legend(self, handles: Iterable[Artist], labels: Iterable[str], **kwargs) -> Legend: ...
+    def legend(self, handles: Iterable[Artist], labels: Iterable[str],
+               *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, *, handles: Iterable[Artist], **kwargs) -> Legend: ...
+    def legend(self, *, handles: Iterable[Artist],
+               loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, labels: Iterable[str], **kwargs) -> Legend: ...
+    def legend(self, labels: Iterable[str],
+               *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
     @overload
-    def legend(self, **kwargs) -> Legend: ...
+    def legend(self, *, loc: LegendLocType | None = ..., **kwargs) -> Legend: ...
 
     def text(
         self,

lib/matplotlib/legend.pyi

@@ -14,12 +14,13 @@ from matplotlib.transforms import (
     BboxBase,
     Transform,
 )
+from matplotlib.typing import ColorType, LegendLocType
 
 
 import pathlib
 from collections.abc import Iterable
 from typing import Any, Literal, overload
-from .typing import ColorType
+
 
 class DraggableLegend(DraggableOffsetBox):
     legend: Legend
@@ -55,7 +56,7 @@ class Legend(Artist):
         handles: Iterable[Artist | tuple[Artist, ...]],
         labels: Iterable[str],
         *,
-        loc: str | tuple[float, float] | int | None = ...,
+        loc: LegendLocType | None = ...,
         numpoints: int | None = ...,
         markerscale: float | None = ...,
         markerfirst: bool = ...,
@@ -118,7 +119,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_loc(self, loc: LegendLocType | None = ...) -> None: ...
     def set_title(
         self, title: str, prop: FontProperties | str | pathlib.Path | None = ...
     ) -> None: ...

lib/matplotlib/typing.py

@@ -150,3 +150,23 @@
     ResizeEventType,
     CloseEventType,
 ]
+
+LegendLocType: TypeAlias = (
+    Literal[
+        # for simplicity, we don't distinguish the between allowed positions for
+        # Axes legend and figure legend. It's still better to limit the allowed
+        # range to the union of both rather than to accept arbitrary strings
+        "upper right", "upper left", "lower left", "lower right",
+        "right", "center left", "center right", "lower center", "upper center",
+        "center",
+        # Axes only
+        "best",
+        # Figure only
+        "outside upper left", "outside upper center", "outside upper right",
+        "outside right upper", "outside right center", "outside right lower",
+        "outside lower right", "outside lower center", "outside lower left",
+        "outside left lower", "outside left center", "outside left upper",
+    ] |
+    tuple[float, float] |
+    int
+)