19

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?

Cory Kramer
  • 114,268
  • 16
  • 167
  • 218
  • 1
    Does this answer your question? [Type hinting / annotation (PEP 484) for numpy.ndarray](https://stackoverflow.com/questions/35673895/type-hinting-annotation-pep-484-for-numpy-ndarray) – SamuelNLP Mar 16 '21 at 14:28
  • 1
    [Numpy Typing](https://numpy.org/doc/stable/reference/typing.html#module-numpy.typing) documentation doesn't say anything about possible dimensions. So maybe you cannot easily currently. However with [Python new type](https://docs.python.org/3/library/typing.html#distinct), you may be able to achieve a custom solution. Depends on what you want to do with the type hints, I guess. – NoDataDumpNoContribution Mar 16 '21 at 14:28
  • There is an answer there that specifies the shape of the numpy array... – SamuelNLP Mar 16 '21 at 14:38
  • @SamuelNLP, that link was asked years ago before the addition of the `numpy` typing module. I've had to add a big disclaimer to the accepted answer because people kept giving it a thumbs down. Stick with a reference to the current `numpy` developments. – hpaulj Mar 16 '21 at 17:06

1 Answers1

20

Update on October 24, 2022

This is still not possible currently, but according to this comment on a Numpy GitHub issue, it will be possible once mypy supports PEP 646. Please see the relevant issue on mypy's GitHub repo. That issue is open at the time of writing.

Python 3.11 has been released today with support for PEP646. Once mypy supports PEP646, users will be able to type-hint the shapes of Numpy arrays.

Older answer

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]
jkr
  • 17,119
  • 2
  • 42
  • 68