Type hint 2D numpy array

user_2296458 0
Viewed 4610

As a follow-up to this question I have a function that will return a 2D numpy.array of fixed columns but variable rows

import numpy.typing as npt

def example() -> npt.ArrayLike:
    data = np.array([[1,2,3],
                     [4,5,6],
                     ...,
                     [x,y,z]])

How can I specifically hint that the returned array will be 3 columns by N (variable) rows?

type-hinting numpy python
1 Answers

It seems like it is not possible to type-hint the shape (or data type) of a numpy.ndarray at this point (September 13, 2022). There are, however, some recent pull requests to numpy working towards this goal.

  • https://github.com/numpy/numpy/pull/17719

    makes the np.ndarray class generic w.r.t. its shape and dtype: np.ndarray[~Shape, ~DType]

    However, an explicit non-goal of that PR is to make runtime-subscriptable aliases for numpy.ndarray. According to that PR, those changes will come in later PRs.

  • https://github.com/numpy/numpy/issues/16544

    Issue discussing typing support for shapes. It is still open at the time of writing.

PEP 646 is related to this and has been accepted into Python 3.11. According to numpy/numpy issue #16544, one will be able to type-hint shape and data type of arrays after type checkers like mypy add support for PEP 646.

This is possible to do with the nptyping package, but that is not part of numpy.

from typing import Any
from nptyping import NDArray

# Nx3 array with Any data type.
NDArray[(Any, 3), Any]
user_56660870
Related