Merge pull request #18050 from BvB93/array-like2

MAINT: Add aliases for commonly used `ArrayLike` objects
numpy · Dec 22, 2020 · 557ed6a · 557ed6a
2 parents e3f288e + 864848e
commit 557ed6a
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 7 deletions.
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
@@ -219,9 +219,25 @@ class _8Bit(_16Bit): ...  # type: ignore[misc]
     _ScalarLike,
     _VoidLike,
 )
-from ._array_like import _SupportsArray, ArrayLike
 from ._shape import _Shape, _ShapeLike
 from ._dtype_like import _SupportsDType, _VoidDTypeLike, DTypeLike
+from ._array_like import (
+    ArrayLike,
+    _ArrayLike,
+    _NestedSequence,
+    _SupportsArray,
+    _ArrayLikeBool,
+    _ArrayLikeUInt,
+    _ArrayLikeInt,
+    _ArrayLikeFloat,
+    _ArrayLikeComplex,
+    _ArrayLikeTD64,
+    _ArrayLikeDT64,
+    _ArrayLikeObject,
+    _ArrayLikeVoid,
+    _ArrayLikeStr,
+    _ArrayLikeBytes,
+)
 
 if __doc__ is not None:
     from ._add_docstring import _docstrings

diff --git a/numpy/typing/_array_like.py b/numpy/typing/_array_like.py
@@ -3,8 +3,22 @@
 import sys
 from typing import Any, overload, Sequence, TYPE_CHECKING, Union, TypeVar
 
-from numpy import ndarray, dtype
-from ._scalars import _ScalarLike
+from numpy import (
+    ndarray,
+    dtype,
+    generic,
+    bool_,
+    unsignedinteger,
+    integer,
+    floating,
+    complexfloating,
+    timedelta64,
+    datetime64,
+    object_,
+    void,
+    str_,
+    bytes_,
+)
 from ._dtype_like import DTypeLike
 
 if sys.version_info >= (3, 8):
@@ -18,6 +32,7 @@
     else:
         HAVE_PROTOCOL = True
 
+_T = TypeVar("_T")
 _DType = TypeVar("_DType", bound="dtype[Any]")
 
 if TYPE_CHECKING or HAVE_PROTOCOL:
@@ -30,6 +45,24 @@ def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
 else:
     _SupportsArray = Any
 
+# TODO: Wait for support for recursive types
+_NestedSequence = 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[
+    _NestedSequence[_SupportsArray[_DType]],
+    _NestedSequence[_T],
+]
+
 # TODO: support buffer protocols once
 #
 # https://bugs.python.org/issue27501
@@ -38,8 +71,48 @@ def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
 #
 # https://github.com/python/typing/issues/593
 ArrayLike = Union[
-    _ScalarLike,
-    Sequence[_ScalarLike],
-    Sequence[Sequence[Any]],  # TODO: Wait for support for recursive types
-    "_SupportsArray[Any]",
+    _RecursiveSequence,
+    _ArrayLike[
+        "dtype[Any]",
+        Union[bool, int, float, complex, str, bytes]
+    ],
+]
+
+# `ArrayLike<X>`: array-like objects that can be coerced into `X`
+# given the casting rules `same_kind`
+_ArrayLikeBool = _ArrayLike[
+    "dtype[bool_]",
+    bool,
+]
+_ArrayLikeUInt = _ArrayLike[
+    "dtype[Union[bool_, unsignedinteger[Any]]]",
+    bool,
+]
+_ArrayLikeInt = _ArrayLike[
+    "dtype[Union[bool_, integer[Any]]]",
+    Union[bool, int],
+]
+_ArrayLikeFloat = _ArrayLike[
+    "dtype[Union[bool_, integer[Any], floating[Any]]]",
+    Union[bool, int, float],
+]
+_ArrayLikeComplex = _ArrayLike[
+    "dtype[Union[bool_, integer[Any], floating[Any], complexfloating[Any, Any]]]",
+    Union[bool, int, float, complex],
+]
+_ArrayLikeTD64 = _ArrayLike[
+    "dtype[Union[bool_, integer[Any], timedelta64]]",
+    Union[bool, int],
+]
+_ArrayLikeDT64 = _NestedSequence[_SupportsArray["dtype[datetime64]"]]
+_ArrayLikeObject = _NestedSequence[_SupportsArray["dtype[object_]"]]
+
+_ArrayLikeVoid = _NestedSequence[_SupportsArray["dtype[void]"]]
+_ArrayLikeStr = _ArrayLike[
+    "dtype[str_]",
+    str,
+]
+_ArrayLikeBytes = _ArrayLike[
+    "dtype[bytes_]",
+    bytes,
 ]