summaryrefslogtreecommitdiff
path: root/numpy/typing
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/typing')
-rw-r--r--numpy/typing/__init__.py3
-rw-r--r--numpy/typing/_array_like.py27
-rw-r--r--numpy/typing/_dtype_like.py46
-rw-r--r--numpy/typing/_shape.py6
4 files changed, 82 insertions, 0 deletions
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
new file mode 100644
index 000000000..94f76a91f
--- /dev/null
+++ b/numpy/typing/__init__.py
@@ -0,0 +1,3 @@
+from ._array_like import _SupportsArray, ArrayLike
+from ._shape import _Shape, _ShapeLike
+from ._dtype_like import DtypeLike
diff --git a/numpy/typing/_array_like.py b/numpy/typing/_array_like.py
new file mode 100644
index 000000000..54a612fb4
--- /dev/null
+++ b/numpy/typing/_array_like.py
@@ -0,0 +1,27 @@
+import sys
+from typing import Any, overload, Sequence, Tuple, Union
+
+from numpy import ndarray
+from ._dtype_like import DtypeLike
+
+if sys.version_info >= (3, 8):
+ from typing import Protocol
+ HAVE_PROTOCOL = True
+else:
+ try:
+ from typing_extensions import Protocol
+ except ImportError:
+ HAVE_PROTOCOL = False
+ else:
+ HAVE_PROTOCOL = True
+
+if HAVE_PROTOCOL:
+ class _SupportsArray(Protocol):
+ @overload
+ def __array__(self, __dtype: DtypeLike = ...) -> ndarray: ...
+ @overload
+ def __array__(self, dtype: DtypeLike = ...) -> ndarray: ...
+else:
+ _SupportsArray = Any
+
+ArrayLike = Union[bool, int, float, complex, _SupportsArray, Sequence]
diff --git a/numpy/typing/_dtype_like.py b/numpy/typing/_dtype_like.py
new file mode 100644
index 000000000..b9df0af04
--- /dev/null
+++ b/numpy/typing/_dtype_like.py
@@ -0,0 +1,46 @@
+from typing import Any, Dict, List, Sequence, Tuple, Union
+
+from numpy import dtype
+from ._shape import _ShapeLike
+
+_DtypeLikeNested = Any # TODO: wait for support for recursive types
+
+# Anything that can be coerced into numpy.dtype.
+# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
+DtypeLike = Union[
+ dtype,
+ # default data type (float64)
+ None,
+ # array-scalar types and generic types
+ type, # TODO: enumerate these when we add type hints for numpy scalars
+ # TODO: add a protocol for anything with a dtype attribute
+ # character codes, type strings or comma-separated fields, e.g., 'float64'
+ str,
+ # (flexible_dtype, itemsize)
+ Tuple[_DtypeLikeNested, int],
+ # (fixed_dtype, shape)
+ Tuple[_DtypeLikeNested, _ShapeLike],
+ # [(field_name, field_dtype, field_shape), ...]
+ #
+ # The type here is quite broad because NumPy accepts quite a wide
+ # range of inputs inside the list; see the tests for some
+ # examples.
+ List[Any],
+ # {'names': ..., 'formats': ..., 'offsets': ..., 'titles': ...,
+ # 'itemsize': ...}
+ # TODO: use TypedDict when/if it's officially supported
+ Dict[
+ str,
+ Union[
+ Sequence[str], # names
+ Sequence[_DtypeLikeNested], # formats
+ Sequence[int], # offsets
+ Sequence[Union[bytes, str, None]], # titles
+ int, # itemsize
+ ],
+ ],
+ # {'field1': ..., 'field2': ..., ...}
+ Dict[str, Tuple[_DtypeLikeNested, int]],
+ # (base_dtype, new_dtype)
+ Tuple[_DtypeLikeNested, _DtypeLikeNested],
+]
diff --git a/numpy/typing/_shape.py b/numpy/typing/_shape.py
new file mode 100644
index 000000000..4629046ea
--- /dev/null
+++ b/numpy/typing/_shape.py
@@ -0,0 +1,6 @@
+from typing import Sequence, Tuple, Union
+
+_Shape = Tuple[int, ...]
+
+# Anything that can be coerced to a shape tuple
+_ShapeLike = Union[int, Sequence[int]]