diff options
| -rw-r--r-- | numpy/tests/typing/fail/dtype.py | 15 | ||||
| -rw-r--r-- | numpy/tests/typing/pass/dtype.py | 32 | ||||
| -rw-r--r-- | numpy/tests/typing/pass/simple.py | 2 | ||||
| -rw-r--r-- | numpy/typing/_dtype_like.py | 59 |
4 files changed, 91 insertions, 17 deletions
diff --git a/numpy/tests/typing/fail/dtype.py b/numpy/tests/typing/fail/dtype.py new file mode 100644 index 000000000..3dc027daf --- /dev/null +++ b/numpy/tests/typing/fail/dtype.py @@ -0,0 +1,15 @@ +import numpy as np + + +class Test: + not_dtype = float + + +np.dtype(Test()) # E: Argument 1 to "dtype" has incompatible type + +np.dtype( + { # E: Argument 1 to "dtype" has incompatible type + "field1": (float, 1), + "field2": (int, 3), + } +) diff --git a/numpy/tests/typing/pass/dtype.py b/numpy/tests/typing/pass/dtype.py index f954fdd44..cbae8c078 100644 --- a/numpy/tests/typing/pass/dtype.py +++ b/numpy/tests/typing/pass/dtype.py @@ -1,3 +1,35 @@ import numpy as np np.dtype(dtype=np.int64) +np.dtype(int) +np.dtype("int") +np.dtype(None) + +np.dtype((int, 2)) +np.dtype((int, (1,))) + +np.dtype({"names": ["a", "b"], "formats": [int, float]}) +np.dtype({"names": ["a"], "formats": [int], "titles": [object]}) +np.dtype({"names": ["a"], "formats": [int], "titles": [object()]}) + +np.dtype([("name", np.unicode_, 16), ("grades", np.float64, (2,)), ("age", "int32")]) + +np.dtype( + { + "names": ["a", "b"], + "formats": [int, float], + "itemsize": 9, + "aligned": False, + "titles": ["x", "y"], + "offsets": [0, 1], + } +) + +np.dtype((np.float_, float)) + + +class Test: + dtype = float + + +np.dtype(Test()) diff --git a/numpy/tests/typing/pass/simple.py b/numpy/tests/typing/pass/simple.py index 1fe6df54f..527050557 100644 --- a/numpy/tests/typing/pass/simple.py +++ b/numpy/tests/typing/pass/simple.py @@ -62,8 +62,6 @@ np.dtype(shape_like_dtype) object_dtype = [("field1", object)] np.dtype(object_dtype) -np.dtype({"col1": ("U10", 0), "col2": ("float32", 10)}) -np.dtype((np.int32, {"real": (np.int16, 0), "imag": (np.int16, 2)})) np.dtype((np.int32, (np.int8, 4))) # Dtype comparision diff --git a/numpy/typing/_dtype_like.py b/numpy/typing/_dtype_like.py index b9df0af04..7c1946a3e 100644 --- a/numpy/typing/_dtype_like.py +++ b/numpy/typing/_dtype_like.py @@ -1,10 +1,43 @@ -from typing import Any, Dict, List, Sequence, Tuple, Union +import sys +from typing import Any, List, Sequence, Tuple, Union, TYPE_CHECKING from numpy import dtype from ._shape import _ShapeLike +if sys.version_info >= (3, 8): + from typing import Protocol, TypedDict + HAVE_PROTOCOL = True +else: + try: + from typing_extensions import Protocol, TypedDict + except ImportError: + HAVE_PROTOCOL = False + else: + HAVE_PROTOCOL = True + _DtypeLikeNested = Any # TODO: wait for support for recursive types +if TYPE_CHECKING or HAVE_PROTOCOL: + # Mandatory keys + class _DtypeDictBase(TypedDict): + names: Sequence[str] + formats: Sequence[_DtypeLikeNested] + + # Mandatory + optional keys + class _DtypeDict(_DtypeDictBase, total=False): + offsets: Sequence[int] + titles: Sequence[Any] # Only `str` elements are usable as indexing aliases, but all objects are legal + itemsize: int + aligned: bool + + # A protocol for anything with the dtype attribute + class _SupportsDtype(Protocol): + dtype: _DtypeLikeNested + +else: + _DtypeDict = Any + _SupportsDtype = Any + # Anything that can be coerced into numpy.dtype. # Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html DtypeLike = Union[ @@ -13,7 +46,8 @@ DtypeLike = Union[ 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 + # anything with a dtype attribute + _SupportsDtype, # character codes, type strings or comma-separated fields, e.g., 'float64' str, # (flexible_dtype, itemsize) @@ -28,19 +62,14 @@ DtypeLike = Union[ 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]], + _DtypeDict, # (base_dtype, new_dtype) Tuple[_DtypeLikeNested, _DtypeLikeNested], ] + +# NOTE: while it is possible to provide the dtype as a dict of +# dtype-like objects (e.g. `{'field1': ..., 'field2': ..., ...}`), +# this syntax is officially discourged and +# therefore not included in the Union defining `DtypeLike`. +# +# See https://github.com/numpy/numpy/issues/16891 for more details. |