Thank you for the quick review!
Adding variables as arguments is not something I typically do, so it was a complete oversight on my part.
I would have preferred the Axes to be inferred when indexing, but I also understand the limitation and previous discussions, so I also think a scaled back version is reasonable.

Note that this might also be better served by #25937.

The motivation for this PR is that I noticed a loss in development experience when using mpl>=3.8 compared to mpl<3.8, which benefited from microsoft/python-type-stubs.
I think it would be great to define a class better than ndarray, but it would be painful to develop with object-oriented-interface fig, ax = plt.subplots() until it it implemented.

Looks like Microsoft did: https://github.com/microsoft/python-type-stubs/blob/e328827329a072baac06e287bd6146ff80796e5e/stubs/matplotlib/pyplot.pyi#L114

So the microsoft ones contain 9 different (but mostly redundant by my analysis) overloaded signatures:

A) squeeze=false, nrows=ncols=int -> ndarray
   - One of the cases in this PR
    
B) squeeze=True, nrows=ncols=1 -> Axes
   - One of the cases in this PR
            
C) Squeeze=True, nrows=1 ncols=int -> ndarray
   - Technically wrong, to my understanding, as ncols can still be 1 but only known to be int at type check time
         
D) Squeeze=True, nrows=int ncols=1 -> ndarray
   - Technically wrong, to my understanding, for the same reason
    
E) Squeeze=True, nrows=ncols=int; keyword only -> ndarray
   - Technically wrong, to my understanding
   - redundant to nrows/ncols not being kwonly  (A/I)

F) squeeze=True, nrows omitted, ncols=int; kwonly -> ndarray
   - redundant to C
  
G) squeeze=True, nrows=int, ncols omitted -> ndarray
   - redundant to D

H) squeeze=True, nrows and ncols omitted -> Axes
   - redundant to B
    
I) squeeze=bool, nrows=ncols-int -> ndarray
   - Technically wrong, as squeeze=True, nrows=1,ncols=1 will match but not return an ndarray
   - Closest to the last one from this PR (though there is a union to ensure technically correct)

I think what is in this PR is better than what Microsoft had because it reduces redundant cases, and does not include incorrect overlapping cases (which I'm a little surprised doesn't flag for them, to be honest, I've definitely seen similar errors flagged by mypy) but still at least carves out where we can tell based on input types and the most common calls.

Additionally, the microsoft ones are missing the width/height_ratios arguments and do not specify that the dictionaries have string keys.

I would like to see the three cases from this PR also included in figure.pyi and gridspec.pyi versions of subplots. (These already have 2 overloads, but could be made better by including the Literal[1] case)

I will note that they have a SubplotBase as a possible return, which should also be removed. It doesn't seem to fail mypy without it, and seems to come back to Figure.add_subplot's docstring which stated (but has now removed) it was a possible return, but then it was omitted from the type hints with a comment:

there are a lot of links and comments.
please advise if the original request is complete...

I would like to see the overloads unified to the three cases from this PR in all three places where subplots method exists:

• matplotlib/lib/matplotlib/figure.pyi

  Line 95 in 955432c

  def subplots(

  • Has 2 overloads, but should include the three from this PR
  • Remove SubplotBase, as that is actually outdated

• matplotlib/lib/matplotlib/gridspec.pyi

  Line 41 in 955432c

  def subplots(

  • Same conditions as figure version

GridSpecBase.subplots has no arguments for narrowing :(

This PR does not work as intended. Using the following code from here:

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 2 * np.pi, 400)
y = np.sin(x ** 2)
fig, axs = plt.subplots(2)  # or nrows=2
fig.suptitle('Vertically stacked subplots')
axs[0].plot(x, y)
axs[1].plot(x, -y)

mypy complains with:

> mypy test.py
test.py:8: error: Value of type "Axes | ndarray[Any, Any]" is not indexable  [index]
test.py:9: error: Value of type "Axes | ndarray[Any, Any]" is not indexable  [index]
Found 2 errors in 1 file (checked 1 source file)

There seems to be an issue with the overload. If we can figure that out, great. If not, we will have to revert back to Any and treat this function as being dynamic (static type analysis is not possible).

How did I notice this? We use multiple subplots quite heavily in TorchGeo. The PR that updated matplotlib from 3.9.0 to 3.9.1 results in 400+ mypy errors as a result of this PR. We could add type hints to all 400+ locations to clarify that it is indeed a numpy array, but given that we're talking about 400+ locations, I would much rather fix this in 1 location in matplotlib instead.

P.S. Let me know if you want me to open a formal issue for this. I thought it would be easier to get the attention of everyone involved in this PR by commenting directly on the PR, but I realize it's easier to track bug reports by having a formal issue opened.

Possible solution: explicitly passing squeeze=False makes type checking much easier. I would also be okay with using a return type of np.ndarray for squeeze=False and Any for squeeze=True if we can't get nrows/ncols working properly with static type checking.

Please review my fix in #28518, it passes all of my tests but more eyes are always better.