Bug summary

The type stubs for Figure.subplots has two @overloads. One is generic, and the other is specialised for squeeze=False (in which case the return value is always nd.ndarray). But there are also some other cases where the type could be more accurately inferred.

Code for reproduction

from matplotlib.figure import Figure

fig = Figure()
reveal_type(fig.subplots())
reveal_type(fig.subplots(2))
reveal_type(fig.subplots(1, 2))

Actual outcome

When run in mypy:

infer.py:4: note: Revealed type is "Union[numpy.ndarray[Any, Any], matplotlib.axes.SubplotBase, matplotlib.axes._axes.Axes]"
infer.py:5: note: Revealed type is "Union[numpy.ndarray[Any, Any], matplotlib.axes.SubplotBase, matplotlib.axes._axes.Axes]"
infer.py:6: note: Revealed type is "Union[numpy.ndarray[Any, Any], matplotlib.axes.SubplotBase, matplotlib.axes._axes.Axes]"

Expected outcome

infer.py:4: note: Revealed type is "Union[matplotlib.axes.SubplotBase, matplotlib.axes._axes.Axes]"
infer.py:5: note: Revealed type is "numpy.ndarray[Any, Any]]"
infer.py:6: note: Revealed type is "numpy.ndarray[Any, Any]]"

Additional information

The last two examples may be impossible to implement, because you'd need a way to type "a literal integer != 1". But the first case is the one I initially ran into. I was able to work around it by replacing it with fig.add_subplot(), but that isn't semantically identical.

Operating system

Ubuntu

Matplotlib Version

3.8.0

Matplotlib Backend

TkAgg

Python version

3.10.12

Jupyter version

No response

Installation

pip