Optional NumPy array not narrowed in max key #16416

Bug Report

Here's an example where I'm allowed to index a NumPy array after narrowing its otherwise optional type when I do so directly, but not when I use it as part of a key for max .

To Reproduce

https://gist.github.com/mypy-play/beca56edcf56c561ebfb37147d31c2ed

That is,

def f(a: Optional[npt.NDArray[np.int64]]):
    if a is None:
        a = np.repeat(1, 5).astype(np.int64)
    print(a[0])
    print(max([0], key=lambda x: a[x]))

Expected Behavior

I wouldn't have expected a[x] to be problematic, since we know a is not None at this point. MyPy itself seems to be aware since the preceding line, containing an a[0], does not cause any errors.

Actual Behavior

main.py:11: error: Value of type "Any | None" is not indexable [index]

Your Environment

• Mypy version used: 1.6.1
• Mypy command-line flags: None.
• Mypy configuration options from mypy.ini (and other config files): None
• Python version used: 3.11.5