Type hinting / annotation (PEP 484) for numpy.ndarray

Question

Has anyone implemented type hinting for the specific numpy.ndarray class?

Right now, I\'m using typing.Any, but it would be nice to have something more specific.

Answer 1

At my company we've been using:

from typing import TypeVar, Generic, Tuple, Union, Optional
import numpy as np

Shape = TypeVar("Shape")
DType = TypeVar("DType")


class Array(np.ndarray, Generic[Shape, DType]):
    """
    Use this to type-annotate numpy arrays, e.g.

        def transform_image(image: Array['H,W,3', np.uint8], ...):
            ...

    """
    pass


def func(arr: Array['N,2', int]):
    return arr*2


print(func(arr = np.array([(1, 2), (3, 4)])))

We actually have a MyPy checker around this that checks that the shapes work out (which we should release at some point). Only thing is it doesn't make PyCharm happy (ie you still get the nasty warning lines):