Specific type annotation for NumPy ndarray using mypy

Question

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.

Does this answer your question? Type hinting / annotation (PEP 484) for numpy.ndarray — mkrieger1
– mkrieger1, Commented Feb 24, 2021 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
– John, Commented Feb 24, 2021 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
– MisterMiyagi, Commented Feb 24, 2021 at 10:51
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]: github.com/numpy/numpy/blob/main/numpy/__init__.pyi — Josiah
– Josiah, Commented Mar 6, 2021 at 18:19

tgpfeiffer · Accepted Answer · 2021-08-17 12:19:49Z

43

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.

answered Aug 17, 2021 at 12:19

tgpfeiffer

1,9082 gold badges21 silver badges23 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Specific type annotation for NumPy ndarray using mypy

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related