Numpy type hints in Python (PEP 484)

Question

I would like to add type hints to a method that takes a numpy array as an input, and returns a string. This numpy array contains floats so I tried:

import numpy as np
def foo(array: np.ndarray[np.float64]) -> str:

But it will not work due to a TypeError: 'type' object is not subscriptable.

I found this but could not follow the discussions!

Etymology

It looks like that link is an experimental package that is still in development. Have you looked at the built-in typing library? – Footworn

#35674395 is an answer I gave several years back. 484 type hints were experimental then, and may still be. And #38006133 – Ladew

There is now an open issue in the numpy github repository regarding type hinting / annotation for numpy types. – Tweak

Answer 1

Check out nptyping. It offers type hints for numpy arrays.

In your case, you would end up with:

from nptyping import NDArray, Float64

def foo(array: NDArray[Float64]) -> str:
    ...

You can check your instances as well:

import numpy as np
from nptyping import NDArray, Float64

arr = np.array([[1.0, 2.0],
                [3.0, 4.0],
                [5.0, 6.0]])

isinstance(arr, NDArray[(3, 2), Float64])  # True.

# Or if you don't want to check the dimensions and their sizes:
isinstance(arr, NDArray[Float64])  # Also True.

Rockefeller

Project status? – Psychogenesis

It is alive. There haven't been any feature requests or bug reports recently though. – Rockefeller

Would be great to also have visibility into the ndarray methods such as argmax, shape , .. – Speller

Support/connection to mypy would be great in this project, i.e. so that if I wrote arr: NDArray[(3, 1), Any] = np.array([[1.0, 2.0], [3.0, 4.0]]) I would get type error through mypy. – Ventilation

Similarly, PyCharm has trouble connecting NDArray to np.ndarray – Chaplain

This doesn't work with pycharm – Milicent

Answer 2

As I explained in this question, there is a little trick with wrapping numpy types with ' :

import numpy as np
from typing import Any

def foo(array: 'np.ndarray[Any , np.dtype[np.float64]]') -> str:
    ...

This works correctly for Pylance type checker on VSCode, and gives no error while execute.

Note that np.ndarray expects 2 arguments, the first one is the shape of the array, the second the dtype of the data inside the array. More info in the doc.

This method is also better than using nptyping since it works for all types, not just the ones described by nptyping. As of now, for instance, there is no support for tuple. see related issue.

Here is another interesting example I gave

Superhighway