27

Support for type annotations was added in NumPy 1.20. I'm trying to figure out how to tell mypy that an array is filled with elements of a particular type, the annotation np.ndarray[np.dcomplex] gives the mypy error "ndarray" expects no type arguments, but 1 given.

EDIT: This question is different from Type hinting / annotation (PEP 484) for numpy.ndarray as that question was asked 4 years ago when there wasn't any official support for type hinting. I'm asking for what is the official way to do this, now that type hinting is actually natively supported by numpy 1.20. The documentation at https://numpy.org/doc/stable/reference/typing.html#module-numpy.typing that the top answer there points towards only seems to say things you shouldn't do with type hinting instead of explaining what you should be doing.

John
  • 723
  • 1
  • 5
  • 12
  • Does this answer your question? [Type hinting / annotation (PEP 484) for numpy.ndarray](https://stackoverflow.com/questions/35673895/type-hinting-annotation-pep-484-for-numpy-ndarray) – mkrieger1 Feb 24 '21 at 10:45
  • That answer points towards the documentation which does not explain the thing that I want to know. The documentation really only seems to say stuff you shouldn't do, but not what you should do. – John Feb 24 '21 at 10:47
  • Please be aware that requests for offsite resources are off-topic. If you are asking *specifically* about how to annotate as an array-of-dtype, there is no need to ask for "someone point to a tutorial"; you might want to [edit] your question to remove this request. – MisterMiyagi Feb 24 '21 at 10:51
  • 1
    I think the short answer is you're just slightly too early. Certainly [the source][1] has references to a generic `ndarray`. However looking at the comparable file in a fresh pip install of up-to-date numpy (1.20.1) it's missing. Realistically your options are wait for it to get properly released, or build your own version from master. In either case, you'll then want to annotate `np.ndarray[Any, np.dcomplex]` because the first Type argument is actually for shape. [1]: https://github.com/numpy/numpy/blob/main/numpy/__init__.pyi – Josiah Mar 06 '21 at 18:19

1 Answers1

29

What you are looking for is the numpy.typing.NDArray class: https://numpy.org/doc/stable/reference/typing.html#numpy.typing.NDArray

numpy.typing.NDArray[A] is an alias for numpy.ndarray[Any, numpy.dtype[A]]:

import numpy as np
import numpy.typing as npt

a: npt.NDArray[np.complex64] = np.zeros((3, 3), dtype=np.complex64)
# reveal_type(a)  # -> numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]]]
print(a)

prints

[[0.+0.j 0.+0.j 0.+0.j]
 [0.+0.j 0.+0.j 0.+0.j]
 [0.+0.j 0.+0.j 0.+0.j]]

Note that even though you annotate a as npt.NDArray[np.complex64], you still need to make sure that you pass the matching dtype to the factory on the right side.

a: npt.NDArray[np.complex64] = np.zeros((3, 3), dtype=np.float32)

passes the mypy check just as well.

tgpfeiffer
  • 1,698
  • 2
  • 18
  • 22