numpy · charris · Sep 20, 2021 · Jan 12, 2021 · Jan 12, 2021 · Jan 12, 2021
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
diff --git a/numpy/core/multiarray.pyi b/numpy/core/multiarray.pyi
@@ -65,7 +65,7 @@ from numpy.typing import (
     NDArray,
     ArrayLike,
     _SupportsArray,
-    _NestedSequence,
+    _FiniteNestedSequence,
     _ArrayLikeBool_co,
     _ArrayLikeUInt_co,
     _ArrayLikeInt_co,
@@ -91,7 +91,7 @@ _DTypeLike = Union[
     Type[_SCT],
     _SupportsDType[dtype[_SCT]],
 ]
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 # Valid time units
 _UnitKind = L[

diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi
@@ -1,12 +1,12 @@
 from typing import TypeVar, overload, List, Sequence, Any, SupportsIndex
 
 from numpy import generic, dtype
-from numpy.typing import ArrayLike, NDArray, _NestedSequence, _SupportsArray
+from numpy.typing import ArrayLike, NDArray, _FiniteNestedSequence, _SupportsArray
 
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
 
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 

diff --git a/numpy/lib/arraypad.pyi b/numpy/lib/arraypad.pyi
@@ -15,7 +15,7 @@ from numpy.typing import (
     ArrayLike,
     NDArray,
     _ArrayLikeInt,
-    _NestedSequence,
+    _FiniteNestedSequence,
     _SupportsArray,
 )
 
@@ -45,7 +45,7 @@ _ModeKind = L[
     "empty",
 ]
 
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 

diff --git a/numpy/lib/index_tricks.pyi b/numpy/lib/index_tricks.pyi
@@ -33,7 +33,7 @@ from numpy.typing import (
     # Arrays
     ArrayLike,
     _NestedSequence,
-    _RecursiveSequence,
+    _FiniteNestedSequence,
     NDArray,
     _ArrayLikeInt,
 
@@ -59,21 +59,19 @@ _ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
 __all__: List[str]
 
 @overload
-def ix_(*args: _NestedSequence[_SupportsDType[_DType]]) -> Tuple[ndarray[Any, _DType], ...]: ...
+def ix_(*args: _FiniteNestedSequence[_SupportsDType[_DType]]) -> Tuple[ndarray[Any, _DType], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[str]) -> Tuple[NDArray[str_], ...]: ...
+def ix_(*args: str | _NestedSequence[str]) -> Tuple[NDArray[str_], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[bytes]) -> Tuple[NDArray[bytes_], ...]: ...
+def ix_(*args: bytes | _NestedSequence[bytes]) -> Tuple[NDArray[bytes_], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[bool]) -> Tuple[NDArray[bool_], ...]: ...
+def ix_(*args: bool | _NestedSequence[bool]) -> Tuple[NDArray[bool_], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[int]) -> Tuple[NDArray[int_], ...]: ...
+def ix_(*args: int | _NestedSequence[int]) -> Tuple[NDArray[int_], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[float]) -> Tuple[NDArray[float_], ...]: ...
+def ix_(*args: float | _NestedSequence[float]) -> Tuple[NDArray[float_], ...]: ...
 @overload
-def ix_(*args: _NestedSequence[complex]) -> Tuple[NDArray[complex_], ...]: ...
-@overload
-def ix_(*args: _RecursiveSequence) -> Tuple[NDArray[Any], ...]: ...
+def ix_(*args: complex | _NestedSequence[complex]) -> Tuple[NDArray[complex_], ...]: ...
 
 class nd_grid(Generic[_BoolType]):
     sparse: _BoolType

diff --git a/numpy/lib/shape_base.pyi b/numpy/lib/shape_base.pyi
@@ -17,7 +17,7 @@ from numpy.typing import (
     ArrayLike,
     NDArray,
     _ShapeLike,
-     _NestedSequence,
+     _FiniteNestedSequence,
      _SupportsDType,
      _ArrayLikeBool_co,
      _ArrayLikeUInt_co,
@@ -31,7 +31,7 @@ from numpy.core.shape_base import vstack
 
 _SCT = TypeVar("_SCT", bound=generic)
 
-_ArrayLike = _NestedSequence[_SupportsDType[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsDType[dtype[_SCT]]]
 
 # The signatures of `__array_wrap__` and `__array_prepare__` are the same;
 # give them unique names for the sake of clarity

diff --git a/numpy/lib/stride_tricks.pyi b/numpy/lib/stride_tricks.pyi
@@ -6,12 +6,12 @@ from numpy.typing import (
     ArrayLike,
     _ShapeLike,
     _Shape,
-    _NestedSequence,
+    _FiniteNestedSequence,
     _SupportsArray,
 )
 
 _SCT = TypeVar("_SCT", bound=generic)
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 

diff --git a/numpy/lib/twodim_base.pyi b/numpy/lib/twodim_base.pyi
@@ -33,7 +33,7 @@ from numpy.typing import (
     _SupportsDType,
     ArrayLike,
     NDArray,
-    _NestedSequence,
+    _FiniteNestedSequence,
     _SupportsArray,
     _ArrayLikeInt_co,
     _ArrayLikeFloat_co,
@@ -55,7 +55,7 @@ _DTypeLike = Union[
     dtype[_SCT],
     _SupportsDType[dtype[_SCT]],
 ]
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 

diff --git a/numpy/lib/type_check.pyi b/numpy/lib/type_check.pyi
@@ -28,7 +28,7 @@ from numpy.typing import (
     _64Bit,
     _SupportsDType,
     _ScalarLike_co,
-    _NestedSequence,
+    _FiniteNestedSequence,
     _SupportsArray,
     _DTypeLikeComplex,
 )
@@ -39,7 +39,7 @@ _SCT = TypeVar("_SCT", bound=generic)
 _NBit1 = TypeVar("_NBit1", bound=NBitBase)
 _NBit2 = TypeVar("_NBit2", bound=NBitBase)
 
-_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 
 class _SupportsReal(Protocol[_T_co]):
     @property

diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
@@ -219,6 +219,8 @@ class _32Bit(_64Bit): ...  # type: ignore[misc]
 class _16Bit(_32Bit): ...  # type: ignore[misc]
 class _8Bit(_16Bit): ...  # type: ignore[misc]
 
+
+from ._nested_sequence import _NestedSequence
 from ._nbit import (
     _NBitByte,
     _NBitShort,
@@ -305,8 +307,7 @@ class _8Bit(_16Bit): ...  # type: ignore[misc]
 from ._array_like import (
     ArrayLike as ArrayLike,
     _ArrayLike,
-    _NestedSequence,
-    _RecursiveSequence,
+    _FiniteNestedSequence,
     _SupportsArray,
     _ArrayLikeInt,
     _ArrayLikeBool_co,

diff --git a/numpy/typing/_array_like.py b/numpy/typing/_array_like.py
@@ -18,6 +18,7 @@
     str_,
     bytes_,
 )
+from ._nested_sequence import _NestedSequence
 
 _T = TypeVar("_T")
 _ScalarType = TypeVar("_ScalarType", bound=generic)
@@ -32,21 +33,23 @@
 class _SupportsArray(Protocol[_DType_co]):
     def __array__(self) -> ndarray[Any, _DType_co]: ...
 
-# TODO: Wait for support for recursive types
-_NestedSequence = Union[
+
+# TODO: Wait until mypy supports recursive objects in combination with typevars
+_FiniteNestedSequence = Union[
     _T,
     Sequence[_T],
     Sequence[Sequence[_T]],
     Sequence[Sequence[Sequence[_T]]],
     Sequence[Sequence[Sequence[Sequence[_T]]]],
 ]
-_RecursiveSequence = Sequence[Sequence[Sequence[Sequence[Sequence[Any]]]]]
 
 # A union representing array-like objects; consists of two typevars:
 # One representing types that can be parametrized w.r.t. `np.dtype`
 # and another one for the rest
 _ArrayLike = Union[
+    _SupportsArray[_DType],
     _NestedSequence[_SupportsArray[_DType]],
+    _T,
     _NestedSequence[_T],
 ]
 
@@ -57,12 +60,9 @@ def __array__(self) -> ndarray[Any, _DType_co]: ...
 # is resolved. See also the mypy issue:
 #
 # https://github.com/python/typing/issues/593
-ArrayLike = Union[
-    _RecursiveSequence,
-    _ArrayLike[
-        dtype,
-        Union[bool, int, float, complex, str, bytes]
-    ],
+ArrayLike = _ArrayLike[
+    dtype,
+    Union[bool, int, float, complex, str, bytes],
 ]
 
 # `ArrayLike<X>_co`: array-like objects that can be coerced into `X`
@@ -95,10 +95,19 @@ def __array__(self) -> ndarray[Any, _DType_co]: ...
     "dtype[Union[bool_, integer[Any], timedelta64]]",
     Union[bool, int],
 ]
-_ArrayLikeDT64_co = _NestedSequence[_SupportsArray["dtype[datetime64]"]]
-_ArrayLikeObject_co = _NestedSequence[_SupportsArray["dtype[object_]"]]
+_ArrayLikeDT64_co = Union[
+    _SupportsArray["dtype[datetime64]"],
+    _NestedSequence[_SupportsArray["dtype[datetime64]"]],
+]
+_ArrayLikeObject_co = Union[
+    _SupportsArray["dtype[object_]"],
+    _NestedSequence[_SupportsArray["dtype[object_]"]],
+]
 
-_ArrayLikeVoid_co = _NestedSequence[_SupportsArray["dtype[void]"]]
+_ArrayLikeVoid_co = Union[
+    _SupportsArray["dtype[void]"],
+    _NestedSequence[_SupportsArray["dtype[void]"]],
+]
 _ArrayLikeStr_co = _ArrayLike[
     "dtype[str_]",
     str,

diff --git a/numpy/typing/_nested_sequence.py b/numpy/typing/_nested_sequence.py
@@ -0,0 +1,93 @@
+"""A module containing the `_NestedSequence` protocol."""
+
+from __future__ import annotations
+
+import sys
+from typing import (
+    Any,
+    Iterator,
+    overload,
+    TypeVar,
+    Protocol,
+)
+
+__all__ = ["_NestedSequence"]
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class _NestedSequence(Protocol[_T_co]):
+    """A protocol for representing nested sequences.
+
+    Warning
+    -------
+    `_NestedSequence` currently does not work in combination with typevars,
+    *e.g.* ``def func(a: _NestedSequnce[T]) -> T: ...``.
+
+    See Also
+    --------
+    `collections.abc.Sequence`
+        ABCs for read-only and mutable :term:`sequences`.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> from __future__ import annotations
+
+        >>> from typing import TYPE_CHECKING
+        >>> import numpy as np
+        >>> from numpy.typing import _NestedSequnce
+
+        >>> def get_dtype(seq: _NestedSequnce[float]) -> np.dtype[np.float64]:
+        ...     return np.asarray(seq).dtype
+
+        >>> a = get_dtype([1.0])
+        >>> b = get_dtype([[1.0]])
+        >>> c = get_dtype([[[1.0]]])
+        >>> d = get_dtype([[[[1.0]]]])
+
+        >>> if TYPE_CHECKING:
+        ...     reveal_locals()
+        ...     # note: Revealed local types are:
+        ...     # note:     a: numpy.dtype[numpy.floating[numpy.typing._64Bit]]
+        ...     # note:     b: numpy.dtype[numpy.floating[numpy.typing._64Bit]]
+        ...     # note:     c: numpy.dtype[numpy.floating[numpy.typing._64Bit]]
+        ...     # note:     d: numpy.dtype[numpy.floating[numpy.typing._64Bit]]
+
+    """
+
+    def __len__(self, /) -> int:
+        """Implement ``len(self)``."""
+        raise NotImplementedError
+
+    @overload
+    def __getitem__(self, index: int, /) -> _T_co | _NestedSequence[_T_co]: ...
+    @overload
+    def __getitem__(self, index: slice, /) -> _NestedSequence[_T_co]: ...
+
+    def __getitem__(self, index, /):
+        """Implement ``self[x]``."""
+        raise NotImplementedError
+
+    def __contains__(self, x: object, /) -> bool:
+        """Implement ``x in self``."""
+        raise NotImplementedError
+
+    def __iter__(self, /) -> Iterator[_T_co | _NestedSequence[_T_co]]:
+        """Implement ``iter(self)``."""
+        raise NotImplementedError
+
+    def __reversed__(self, /) -> Iterator[_T_co | _NestedSequence[_T_co]]:
+        """Implement ``reversed(self)``."""
+        raise NotImplementedError
+
+    def count(self, value: Any, /) -> int:
+        """Return the number of occurrences of `value`."""
+        raise NotImplementedError
+
+    def index(
+        self, value: Any, start: int = 0, stop: int = sys.maxsize, /
+    ) -> int:
+        """Return the first index of `value`."""
+        raise NotImplementedError
diff --git a/numpy/typing/tests/data/fail/comparisons.py b/numpy/typing/tests/data/fail/comparisons.py
@@ -22,7 +22,7 @@
 AR_M > AR_m  # E: Unsupported operand types
 
 # Unfortunately `NoReturn` errors are not the most descriptive
-_1 = AR_i > str()  # E: Need type annotation
+_1 = AR_i > str()  # E: No overload variant
 _2 = AR_i > bytes()  # E: Need type annotation
-_3 = str() > AR_M  # E: Need type annotation
+_3 = str() > AR_M  # E: Unsupported operand types
 _4 = bytes() > AR_M  # E: Need type annotation
diff --git a/numpy/typing/tests/data/fail/multiarray.py b/numpy/typing/tests/data/fail/multiarray.py
@@ -42,7 +42,7 @@ def func(a: int) -> None: ...
 
 np.busday_offset("2012", 10)  # E: incompatible type
 
-np.datetime_as_string("2012")  # E: incompatible type
+np.datetime_as_string("2012")  # E: No overload variant
 
 np.compare_chararrays("a", b"a", "==", False)  # E: No overload variant
 

diff --git a/numpy/typing/tests/data/fail/nested_sequence.py b/numpy/typing/tests/data/fail/nested_sequence.py
@@ -0,0 +1,17 @@
+from typing import Sequence, Tuple, List
+import numpy.typing as npt
+
+a: Sequence[float]
+b: List[complex]
+c: Tuple[str, ...]
+d: int
+e: str
+
+def func(a: npt._NestedSequence[int]) -> None:
+    ...
+
+reveal_type(func(a))  # E: incompatible type
+reveal_type(func(b))  # E: incompatible type
+reveal_type(func(c))  # E: incompatible type
+reveal_type(func(d))  # E: incompatible type
+reveal_type(func(e))  # E: incompatible type