5

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?

Alex Waygood
  • 6,304
  • 3
  • 24
  • 46
palapapa
  • 573
  • 2
  • 5
  • 25
  • 1
    Does this answer your question? [Numpy type hints in Python (PEP 484)](https://stackoverflow.com/questions/52839427/numpy-type-hints-in-python-pep-484) – ddejohn Aug 25 '21 at 05:09
  • @blorgon It doesn't behave like the way it's documented. I wanted to have a `ndarray` that had 3 dimensions of any size and dtype of `int8`, so I wrote it as `nptyping.NDArray[(typing.Any, typing.Any, typing.Any), nptyping.Int8]`, but it gave me this error: `Expected class type but received "None"` – palapapa Aug 25 '21 at 07:05

1 Answers1

3

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

Onyr
  • 769
  • 5
  • 21
  • Does the part with `ThreeD` actually work? I was trying to limit the `ndarray` to 3-dimensional only. That's why I wrote `tuple[int, int, int]` but looking back I don't think it would work. – palapapa Jun 16 '22 at 17:18
  • 1
    There is not "real" limitation in one situation or another. Typing is just something useful on the developer perspective, not really at code execution unless you use `assert(isinstance())`. I kept your custom types for clarity but you can concatenate everything if you prefer. Hope it helps. – Onyr Jun 17 '22 at 07:47