Type hint for NumPy ndarray dtype?

Question

I would like a function to include a type hint for NumPy ndarray\'s alongside with its dtype.

With lists, for example, one could do the fol

Answer 1

One informal solution for type documentation is the following:

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've been using this at my company and made a MyPy checker that actually 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):