API: Add npt.NDArray, a runtime-subscriptable alias for np.ndarray

[BvB93]

This PR introduces a runtime-subscriptable alias for np.ndarray[Any, np.dtype[~Scalar]] by the name of npt.NDArray.

Follow up on #17719, which made the initial attempt at introducing aforementioned type alias before being deferred to a later PR.

Motivation for the introduction of this type alias is two-fold:

• It allows one to annotate an array with a given dtype during runtime, e.g. the same role typing.Sequence fulfilled for collections.abc.Sequence prior to python 3.9. While runetime-subsription is already, more or less, possible to with from __future__ import annotations, the latter only helps with annotations, and not the creation of type aliases. All in all npt.NDArray would add a big convenience factor.
• It provides a compact alias for an otherwise rather verbose expression, a convenience that was already extensive utilized in the numpy stubs via its private predecessor (c332c6a). Note this alias cannot be used for custom dtypes that are not based on np.generic (see NEP 42), which will either have to define their own aliases or stick to the complete form.

Future Work

Currently I can think of two other (potential) candidates:

• An alias for np.dtype[~Scalar]. I'm not 100% sure if the big dtype update (NEP 40-42) already has plans on adding a dtype.__class_getitem__ method or not, so adding a subscriptable alias to numpy.typing might be redudant. @seberg do you have any comments this?
• An alias for np.ndarray[~Shape, np.dtype[~Scalar]]. I'd say this would be worthwhile to introduce, once we have shape-typing sorted out (Typing support for shapes #16544), that is.
  There is an open question of what would be the best name for such type-alias. npt.Array perhaps? To distinguish it from the arbitrary shaped npt.NDarray introduced in this PR.

Implementation

Fortunately PEP 585 has provided us with all the tools we need for constructing what are effectively subscriptable wrappers: types.GenericAlias. Aforementioned class is designed such that (nearly) all operations wrap around the underlying type, e.g. np.ndarray behaves virtually the same as types.GenericAlias(np.ndarray, ()).

The only bad news is that GenericAlias is exclusive to python >= 3.9 (and written in C), hence the introduction of a (python-based) backport in this PR. The backport can of course be removed once we've dropped support for 3.8 at some point in the future.

Examples

>>> import numpy as np
>>> import numpy.typing as npt

>>> print(npt.NDArray)
numpy.ndarray[typing.Any, numpy.dtype[~ScalarType]]

>>> print(npt.NDArray[np.float64])
numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]

>>> NDArrayInt = npt.NDArray[np.int_]
>>> a: NDArrayInt = np.arange(10)

>>> def func(a: npt.ArrayLike) -> npt.NDArray[Any]:
...     return np.array(a)

[seberg]

I am happy to just add __class_getitem__, I have not done it yet, simply because it wasn't one of the "big things to worry about". I like the syntax (and it seems very reasonable for typing).

@BvB93 if it helps you, I can add this over the weekend, we ping the mailing list about comments (I doubt anyone will be bothered) and then you have it.

[BvB93]

I am happy to just add __class_getitem__, I have not done it yet, simply because it wasn't one of the "big things to worry about". I like the syntax (and it seems very reasonable for typing).

That would be nice, yes; it's definitely one of the things on my wish lists.
Personally I wouldn't mind waiting for 1.22 either, so no hurries.

[seberg]

@BvB93 what is the typical behviour? Since np.dtype[scalar] would return subclasses of np.dtype, __class_getitem__ should refuse to work on subclasses (e.g. np.dtype[scalar][scalar] makes no sense) right?

EDIT: Or maybe, I could allow it, but raise an error unless: new = DType[something]; assert new != DType and issubclass(new, DType)

[BvB93]

@BvB93 what is the typical behviour? Since np.dtype[scalar] would return subclasses of np.dtype, __class_getitem__ should refuse to work on subclasses (e.g. np.dtype[scalar][scalar] makes no sense) right?

I would actually propose to just stick same semantics as cpython used in PEP 585 for the builtin types (python/cpython#18239). For example, with the likes of builtins.list their pure python implementation would be equivalent the code below:

import types

def __class_getitem__(cls, args):
    return types.GenericAlias(cls, args)

Or a bit more fancy and and some basic error checking (i.e. this will raise for dtype[Any, Any] and dtype[()]):

import types

Scalar = TypeVar("Scalar")
_ALIAS = types.GenericAlias(cls, (Scalar,))

def __class_getitem__(cls, args):
    return _ALIAS[args]  # Let `GenericAlias` handle the error checking

This would have a few consequences though:

• types.GenericAlias is only available for python >= 3.9. So we could either limit __class_getitem__, or for example use the backport introduced in this PR. I'd be ok with either option.
• The return type would, during runtime, be a wrapped version of the original class (so not an actual subclass).

[seberg]

The return type would, during runtime, be a wrapped version of the original class (so not an actual subclass).

Wait, but that sounds pretty much useless at run-time? I thought this would be useful semantics to get the actual subclass.

[shoyer]

No objections from me -- this looks very handy!

[BvB93]

Wait, but that sounds pretty much useless at run-time? I thought this would be useful semantics to get the actual subclass.

Well... yes.

Going the __class_getitem__(cls, args) -> SubClass route would be an option, but I'd imagine it's more work to implement.
For example, we'd need a way of handling:

• Typvars: def func(dt: np.dtype[T]) -> T: ...
• Literal strings: IntDtype: np.dtype["np.int_"]
• Any: np.dtype[Any]

Then again, I'm not opposing either option.

[seberg]

@BvB93 no, its not really any work at all, because I already have that. The point is that this is what happens if you type np.array():

np.array([scalar])

must already be able to effectively do:

DType = np.dtype[type(scalar)]  # EDIT: added missing type.
descriptor = DType.get_descriptor_from_pyobject(scalar)
result = np.empty(1, dtype=descriptor)
result[0] = scalar

In other words: I already have a function that does DType = np.dtype[scalar] in NumPy. (Of course there could be DTypes without scalars in the future, but in that case the user would have to use that DTypes name directly in any case.)

[seberg]

Sorry, I hadn't parsed the typing side of what it means to return subclasses... I am not too fond about strings at runtime, but if typing wants it, sure.

From the runtime side of things, I feel it would be natural to have:

SubClassDType = np.dtype[scalar_type]
SubClassDType.type is scalar_type  # we could also call it `scalar_type` or so...

(both of these are guaranteed. I am not 100% sure about the opposite direction that np.dtype[DType.type] must work, but the above must always work.)

That would allow doing things like:

isinstance(arr.dtype, np.dtype[np.float64])

But, of course we probably don't need that all that often...

[BvB93]

FYI the discussion on dtype.__class_getitem__ from this weeks community meeting seems to have stirred towards the types.GenericAlias-based approach (xref python/cpython#18239). As a reminder, such an enhancement would still be beyond the scope of this particular PR and reserved for future work.

Secondly, I'd like to point the attention of (potential) reviewers to the snippet below, as it would be useful to sort out some of these details already (most importantly: would the names make sense?).

Future Work

• An alias for np.ndarray[~Shape, np.dtype[~Scalar]]. I'd say this would be worthwhile to introduce, once we have shape-typing sorted out (Typing support for shapes #16544), that is.
  There is an open question of what would be the best name for such type-alias. npt.Array perhaps? To distinguish it from the arbitrary shaped npt.NDarray introduced in this PR.

[charris]

Thanks Bas.