ENH: Type the possible str legend locs as Literals

timhoffm · timhoffm · commit f608a46ff272 · 2025-01-12T14:02:50.000+01:00
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.
diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -38,7 +38,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
 
 class Axes(_AxesBase):
     def get_title(self, loc: Literal["left", "center", "right"] = ...) -> str: ...
@@ -60,7 +60,13 @@ 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: ...
     @overload
diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi
@@ -25,7 +25,7 @@ from matplotlib.lines import Line2D
 from matplotlib.patches import Rectangle, Patch
 from matplotlib.text import Text
 from matplotlib.transforms import Affine2D, Bbox, BboxBase, Transform
-from .typing import ColorType, HashableList
+from .typing import ColorType, HashableList, LegendLocType
 
 _T = TypeVar("_T")
 
@@ -147,13 +147,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,
diff --git a/lib/matplotlib/legend.pyi b/lib/matplotlib/legend.pyi
@@ -19,7 +19,7 @@ from matplotlib.transforms import (
 import pathlib
 from collections.abc import Iterable
 from typing import Any, Literal, overload
-from .typing import ColorType
+from .typing import ColorType, LegendLocType
 
 class DraggableLegend(DraggableOffsetBox):
     legend: Legend
@@ -55,7 +55,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 +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_loc(self, loc: LegendLocType | None = ...) -> None: ...
     def set_title(
         self, title: str, prop: FontProperties | str | pathlib.Path | None = ...
     ) -> None: ...
diff --git a/lib/matplotlib/typing.py b/lib/matplotlib/typing.py
@@ -80,3 +80,24 @@
 _HT = TypeVar("_HT", bound=Hashable)
 HashableList: TypeAlias = list[_HT | "HashableList[_HT]"]
 """A nested list of Hashable values."""
+
+
+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
+)
