10000 TYP: Add common type overloads of subplot_mosaic · matplotlib/matplotlib@757bb2f · GitHub
[go: up one dir, main page]
Skip to content

Commit 757bb2f

Browse files
QuLogicQuLogic
committed
TYP: Add common type overloads of subplot_mosaic
I'll assert, without proof, that passing a single string, a mosaic list of strings, or a mosaic list all of the same type is more common than passing arbitrary unrelated hashables. Thus it is somewhat convenient if the return type stipulates that the resulting dictionary is also keyed with strings or the common type. This also fixes the type of the `per_subplot_kw` argument, which also allows dictionary keys of tuples of the entries.
1 parent 3ab6e1f commit 757bb2f

File tree

4 files changed

+106
-27
lines changed

4 files changed

+106
-27
lines changed

‎lib/matplotlib/figure.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1798,15 +1798,11 @@ def _norm_per_subplot_kw(per_subplot_kw):
17981798
if isinstance(k, tuple):
17991799
for sub_key in k:
18001800
if sub_key in expanded:
1801-
raise ValueError(
1802-
f'The key {sub_key!r} appears multiple times.'
1803-
)
1801+
raise ValueError(f'The key {sub_key!r} appears multiple times.')
18041802
expanded[sub_key] = v
18051803
else:
18061804
if k in expanded:
1807-
raise ValueError(
1808-
f'The key {k!r} appears multiple times.'
1809-
)
1805+
raise ValueError(f'The key {k!r} appears multiple times.')
18101806
expanded[k] = v
18111807
return expanded
18121808

‎lib/matplotlib/figure.pyi

Lines changed: 40 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
1+
from collections.abc import Callable, Hashable, Iterable
12
import os
3+
from typing import Any, IO, Literal, TypeVar, overload
4+
5+
import numpy as np
6+
from numpy.typing import ArrayLike
27

38
from matplotlib.artist import Artist
49
from matplotlib.axes import Axes, SubplotBase
@@ -19,14 +24,10 @@ from matplotlib.lines import Line2D
1924
from matplotlib.patches import Rectangle, Patch
2025
from matplotlib.text import Text
2126
from matplotlib.transforms import Affine2D, Bbox, BboxBase, Transform
22-
23-
import numpy as np
24-
from numpy.typing import ArrayLike
25-
26-
from collections.abc import Callable, Iterable
27-
from typing import Any, IO, Literal, overload
2827
from .typing import ColorType, HashableList
2928

29+
_T = TypeVar("_T")
30+
3031
class SubplotParams:
3132
def __init__(
3233
self,
@@ -226,21 +227,48 @@ class FigureBase(Artist):
226227
*,
227228
bbox_extra_artists: Iterable[Artist] | None = ...,
228229
) -> Bbox: ...
229-
230-
# Any in list of list is recursive list[list[Hashable | list[Hashable | ...]]] but that can't really be type checked
230+
@overload
231231
def subplot_mosaic(
232232
self,
233-
mosaic: str | HashableList,
233+
mosaic: str,
234+
*,
235+
sharex: bool = ...,
236+
sharey: bool = ...,
237+
width_ratios: ArrayLike | None = ...,
238+
height_ratios: ArrayLike | None = ...,
239+
empty_sentinel: str = ...,
240+
subplot_kw: dict[str, Any] | None = ...,
241+
per_subplot_kw: dict[str | tuple[str, ...], dict[str, Any]] | None = ...,
242+
gridspec_kw: dict[str, Any] | None = ...,
243+
) -> dict[str, Axes]: ...
244+
@overload
245+
def subplot_mosaic(
246+
self,
247+
mosaic: list[HashableList[_T]],
248+
*,
249+
sharex: bool = ...,
250+
sharey: bool = ...,
251+
width_ratios: ArrayLike | None = ...,
252+
height_ratios: ArrayLike | None = ...,
253+
empty_sentinel: _T = ...,
254+
subplot_kw: dict[str, Any] | None = ...,
255+
per_subplot_kw: dict[_T | tuple[_T, ...], dict[str, Any]] | None = ...,
256+
gridspec_kw: dict[str, Any] | None = ...,
257+
) -> dict[_T, Axes]: ...
258+
@overload
259+
def subplot_mosaic(
260+
self,
261+
mosaic: list[HashableList[Hashable]],
234262
*,
235263
sharex: bool = ...,
236264
sharey: bool = ...,
237265
width_ratios: ArrayLike | None = ...,
238266
height_ratios: ArrayLike | None = ...,
239267
empty_sentinel: Any = ...,
240268
subplot_kw: dict[str, Any] | None = ...,
241-
per_subplot_kw: dict[Any, dict[str, Any]] | None = ...,
242-
gridspec_kw: dict[str, Any] | None = ...
243-
) -> dict[Any, Axes]: ...
269+
per_subplot_kw: dict[Hashable | tuple[Hashable, ...], dict[str, Any]] | None = ...,
270+
gridspec_kw: dict[str, Any] | None = ...,
271+
) -> dict[Hashable, Axes]: ...
244272

245273
class SubFigure(FigureBase):
246274
figure: Figure

‎lib/matplotlib/pyplot.py

Lines changed: 61 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@
125125

126126
_P = ParamSpec('_P')
127127
_R = TypeVar('_R')
128+
_T = TypeVar('_T')
128129

129130

130131
# We may not need the following imports here:
@@ -1600,8 +1601,56 @@ def subplots(
16001601
return fig, axs
16011602

16021603

1604+
@overload
1605+
def subplot_mosaic(
1606+
mosaic: str,
1607+
*,
1608+
sharex: bool = ...,
1609+
sharey: bool = ...,
1610+
width_ratios: ArrayLike | None = ...,
1611+
height_ratios: ArrayLike | None = ...,
1612+
empty_sentinel: str = ...,
1613+
subplot_kw: dict[str, Any] | None = ...,
1614+
gridspec_kw: dict[str, Any] | None = ...,
1615+
per_subplot_kw: dict[str | tuple[str, ...], dict[str, Any]] | None = ...,
1616+
**fig_kw: Any
1617+
) -> tuple[Figure, dict[str, matplotlib.axes.Axes]]: ...
1618+
1619+
1620+
@overload
1621+
def subplot_mosaic(
1622+
mosaic: list[HashableList[_T]],
1623+
*,
1624+
sharex: bool = ...,
1625+
sharey: bool = ...,
1626+
width_ratios: ArrayLike | None = ...,
1627+
height_ratios: ArrayLike | None = ...,
1628+
empty_sentinel: _T = ...,
1629+
subplot_kw: dict[str, Any] | None = ...,
1630+
gridspec_kw: dict[str, Any] | None = ...,
1631+
per_subplot_kw: dict[_T | tuple[_T, ...], dict[str, Any]] | None = ...,
1632+
**fig_kw: Any
1633+
) -> tuple[Figure, dict[_T, matplotlib.axes.Axes]]: ...
1634+
1635+
1636+
@overload
1637+
def subplot_mosaic(
1638+
mosaic: list[HashableList[Hashable]],
1639+
*,
1640+
sharex: bool = ...,
1641+
sharey: bool = ...,
1642+
width_ratios: ArrayLike | None = ...,
1643+
height_ratios: ArrayLike | None = ...,
1644+
empty_sentinel: Any = ...,
1645+
subplot_kw: dict[str, Any] | None = ...,
1646+
gridspec_kw: dict[str, Any] | None = ...,
1647+
per_subplot_kw: dict[Hashable | tuple[Hashable, ...], dict[str, Any]] | None = ...,
1648+
**fig_kw: Any
1649+
) -> tuple[Figure, dict[Hashable, matplotlib.axes.Axes]]: ...
1650+
1651+
16031652
def subplot_mosaic(
1604-
mosaic: str | HashableList,
1653+
mosaic: str | list[HashableList[_T]] | list[HashableList[Hashable]],
16051654
*,
16061655
sharex: bool = False,
16071656
sharey: bool = False,
@@ -1610,9 +1659,13 @@ def subplot_mosaic(
16101659
empty_sentinel: Any = '.',
16111660
subplot_kw: dict[str, Any] | None = None,
16121661
gridspec_kw: dict[str, Any] | None = None,
1613-
per_subplot_kw: dict[Hashable, dict[str, Any]] | None = None,
1614-
**fig_kw
1615-
) -> tuple[Figure, dict[Hashable, matplotlib.axes.Axes]]:
1662+
per_subplot_kw: dict[str | tuple[str, ...], dict[str, Any]] |
1663+
dict[_T | tuple[_T, ...], dict[str, Any]] |
1664+
dict[Hashable | tuple[Hashable, ...], dict[str, Any]] | None = None,
1665+
**fig_kw: Any
1666+
) -> tuple[Figure, dict[str, matplotlib.axes.Axes]] | \
1667+
tuple[Figure, dict[_T, matplotlib.axes.Axes]] | \
1668+
tuple[Figure, dict[Hashable, matplotlib.axes.Axes]]:
16161669
"""
16171670
Build a layout of Axes based on ASCII art or nested lists.
16181671
@@ -1714,12 +1767,13 @@ def subplot_mosaic(
17141767
17151768
"""
17161769
fig = figure(**fig_kw)
1717-
ax_dict = fig.subplot_mosaic(
1718-
mosaic, sharex=sharex, sharey=sharey,
1770+
ax_dict = fig.subplot_mosaic( # type: ignore[misc]
1771+
mosaic, # type: ignore[arg-type]
1772+
sharex=sharex, sharey=sharey,
17191773
height_ratios=height_ratios, width_ratios=width_ratios,
17201774
subplot_kw=subplot_kw, gridspec_kw=gridspec_kw,
17211775
empty_sentinel=empty_sentinel,
1722-
per_subplot_kw=per_subplot_kw,
1776+
per_subplot_kw=per_subplot_kw, # type: ignore[arg-type]
17231777
)
17241778
return fig, ax_dict
17251779

‎lib/matplotlib/typing.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
"""
1212
from collections.abc import Hashable, Sequence
1313
import pathlib
14-
from typing import Any, Literal, Union
14+
from typing import Any, Literal, TypeVar, Union
1515

1616
from . import path
1717
from ._enums import JoinStyle, CapStyle
@@ -55,5 +55,6 @@
5555
Sequence[Union[str, pathlib.Path, dict[str, Any]]],
5656
]
5757

58-
HashableList = list[Union[Hashable, "HashableList"]]
58+
_HT = TypeVar("_HT", bound=Hashable)
59+
HashableList = list[Union[_HT, "HashableList[_HT]"]]
5960
"""A nested list of Hashable values."""

0 commit comments

Comments
 (0)
0