-
-
Notifications
You must be signed in to change notification settings - Fork 11.7k
Description
Proposed new feature or change:
I believe _ShapeType should be covariant. For instance, the following code does not typecheck because _ShapeType is invariant:
from typing import NewType
import numpy as np
from numpy.typing import NBitBase
Time = NewType("Time", int)
Series = NewType("Series", int)
AnyFloat = np.dtype[np.floating[NBitBase]]
arr: np.ndarray[tuple[Time, Series], AnyFloat] = np.arange(4.0).reshape((2, 2))
def foo(a: np.ndarray[tuple[int, ...], AnyFloat]) -> None: ...
foo(arr)I don't think there's any problems making it covariant, because I would assume it will be replaced by TypeVarTuple (PEP 646), and tuples are covariant. In the interim, there's a few functions that change the shape of the array, such as __getitem__, but it seems those return ndarray[Any,] except in the overload:
@overload
def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType, _dtype[void]]: ...I'm not sure what situation this is handling, since I've never seen self typed.
There's also the chararray, memmap, matrix, recordarray, MaskedArray, and MaskedRecord types (of which the latter two define an identical _ShapeType type variable), but their methods also appear to return Any as the shape of any return that isn't the same shape.
Would you mind if I PR'd this and saw if it passes type checking?