From 631e38f1e73cfa82482f80ce8752d315257a754a Mon Sep 17 00:00:00 2001 From: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com> Date: Sun, 12 Jan 2025 13:43:55 +0100 Subject: [PATCH] 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. --- lib/matplotlib/axes/_axes.pyi | 13 ++++++++----- lib/matplotlib/figure.pyi | 13 ++++++++----- lib/matplotlib/legend.pyi | 7 ++++--- lib/matplotlib/typing.py | 21 +++++++++++++++++++++ 4 files changed, 41 insertions(+), 13 deletions(-) diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi index 1877cc192b15..e6fe5bcaf65e 100644 --- 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,13 +60,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, diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi index 08bf1505532b..f6d26eab4633 100644 --- 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 index dde5882da69d..c03471fc54d1 100644 --- a/lib/matplotlib/legend.pyi +++ b/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: ... diff --git a/lib/matplotlib/typing.py b/lib/matplotlib/typing.py index e29b08b437d6..b2cb9f48ebce 100644 --- 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 +)