Mypy is not aware that changing the dtype of an array changes its annotated type, even if the dtype is passed as a literal np.dtype[np.float64], np.dtype[np.float32], etc.

I imagine that this kind of type inference could be very difficult in general, but it would be nice if at least this special case worked.

Also - the parameters of np.ndarray don't appear to be documented anywhere. I kind of just guessed that it was something like tuple[int], where the length of the tuple represented the shape, but honestly I have no idea if this is right. It would be great if this was stated in the docs somewhere!

Reproducing code example:

import numpy as np

Vec64 = np.ndarray[tuple[int], np.dtype[np.float64]]
Vec32 = np.ndarray[tuple[int], np.dtype[np.float32]]

def convert(data: Vec64) -> Vec32:
    result = data.astype(np.float32)
    reveal_type(result)
    return result

x = np.array([1, 2, 3], dtype=np.float64)
y = convert(x)

Error message:

Mypy output:

% mypy mypy_numpy_test.py
mypy_numpy_test.py:8: note: Revealed type is "numpy.ndarray*[Tuple[builtins.int], numpy.dtype[numpy.floating[numpy.typing._64Bit]]]"
mypy_numpy_test.py:9: error: Incompatible return value type (got "ndarray[Tuple[int], dtype[floating[_64Bit]]]", expected "ndarray[Tuple[int], dtype[floating[_32Bit]]]")
Found 1 error in 1 file (checked 1 source file)

NumPy/Python version information:

In [3]: import sys

In [4]: print(sys.version)
3.9.6 (default, Aug  3 2021, 19:43:34)
[Clang 10.0.1 (clang-1001.0.46.4)]

In [5]: import numpy

In [6]: print(numpy.__version__)
1.21.2

In [7]: !mypy --version
mypy 0.910

mypy.ini:

[run]
strict = true

[mypy]
plugins =
	numpy.typing.mypy_plugin,
	classes.contrib.mypy.classes_plugin,
	returns.contrib.mypy.returns_plugin