pythonnumpypython-typingtype-alias

TypeError: 'numpy._DTypeMeta' object is not subscriptable


I'm trying to type hint a numpy ndarray like this:

RGB = numpy.dtype[numpy.uint8]
ThreeD = tuple[int, int, int]

def load_images(paths: list[str]) -> tuple[list[numpy.ndarray[ThreeD, RGB]], list[str]]: ...

but at the first line when I run this, I got the following error:

RGB = numpy.dtype[numpy.uint8]
TypeError: 'numpy._DTypeMeta' object is not subscriptable

How do I type hint a ndarray correctly?


Solution

  • It turns out that strongly type a numpy array is not straightforward at all. I spent a couple of hours to figure out how to do it properly.

    A simple method that do not add yet another dependency to your project is to use a trick described here. Just wrap numpy types with with ':

    import numpy
    import numpy.typing as npt
    from typing import cast, Type, Sequence
    import typing
    
    RGB: typing.TypeAlias = 'numpy.dtype[numpy.uint8]'
    ThreeD: typing.TypeAlias = tuple[int, int, int]
    NDArrayRGB: typing.TypeAlias = 'numpy.ndarray[ThreeD, RGB]'
    
    def load_images(paths: list[str]) -> tuple[list[NDArrayRGB], list[str]]: ...
    

    The trick is to use single-quotes to avoid the infamous TypeError: 'numpy._DTypeMeta' object is not subscriptable when Python tries to interpret the [] in the expression. This trick is well handled for instance by VSCode Pylance type-checker:

    enter image description here

    Notice that the colors for types are respected and that the execution gives no error.

    Note about nptyping

    As suggested by @ddejohn, one can use nptyping. Just install the package: pip install nptyping. However, as of now (16 June 2022), there is no Tuple type defined in nptyping so you won't be able to prefectly type you code that way. I have open a new issue so maybe in the future it will work.

    edits

    Turns out there is a different way to express a tuple as a nptyping.Shape as answered by ramonhagenaars, which is also elegant:

    from nptyping import NDArray, Shape, UInt8
    
    # A 1-dimensional array (i.e. 1 RGB color).
    RGBArray1D = NDArray[Shape["[r, g, b]"], UInt8]
    
    # A 2-dimensional array (i.e. an array of RGB colors).
    RGBArrayND = NDArray[Shape["*, [r, g, b]"], UInt8]
    
    def load_images_trick(paths: list[str]) -> tuple[list[RGBArrayND], list[str]]: ...
    

    However, this solution is not well supported by VSCode Pylance, an I get an error suggestion for Shape:

    Expected class type but received "Literal"
      "Literal" is not a class
      "Literal" is not a classPylancereportGeneralTypeIssues
    Pylance(reportGeneralTypeIssues)
    

    enter image description here