Skip to content

Commit a72a046

Browse files
jorenhamjorenham
authored
Merge pull request #481 from numpy/_numtype.ToDType
✨ `_numtype` dtype-like type-aliases
2 parents ccff50a + 635d8d7 commit a72a046

File tree

5 files changed

+180
-67
lines changed

5 files changed

+180
-67
lines changed

‎src/_numtype/__init__.pyi

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,32 @@ from typing_extensions import Protocol, TypeAliasType, TypeVar, Unpack
1111
import numpy as np
1212
from numpy._typing import _NestedSequence
1313

14+
from ._dtype import (
15+
ToDType as ToDType,
16+
ToDTypeBool as ToDTypeBool,
17+
ToDTypeBytes as ToDTypeBytes,
18+
ToDTypeCLongDouble as ToDTypeCLongDouble,
19+
ToDTypeComplex64 as ToDTypeComplex64,
20+
ToDTypeComplex128 as ToDTypeComplex128,
21+
ToDTypeDateTime64 as ToDTypeDateTime64,
22+
ToDTypeFloat16 as ToDTypeFloat16,
23+
ToDTypeFloat32 as ToDTypeFloat32,
24+
ToDTypeFloat64 as ToDTypeFloat64,
25+
ToDTypeInt8 as ToDTypeInt8,
26+
ToDTypeInt16 as ToDTypeInt16,
27+
ToDTypeInt32 as ToDTypeInt32,
28+
ToDTypeInt64 as ToDTypeInt64,
29+
ToDTypeLongDouble as ToDTypeLongDouble,
30+
ToDTypeObject as ToDTypeObject,
31+
ToDTypeStr as ToDTypeStr,
32+
ToDTypeString as ToDTypeString,
33+
ToDTypeTimeDelta64 as ToDTypeTimeDelta64,
34+
ToDTypeUInt8 as ToDTypeUInt8,
35+
ToDTypeUInt16 as ToDTypeUInt16,
36+
ToDTypeUInt32 as ToDTypeUInt32,
37+
ToDTypeUInt64 as ToDTypeUInt64,
38+
ToDTypeVoid as ToDTypeVoid,
39+
)
1440
from ._just import (
1541
Just as Just,
1642
JustBytes as JustBytes,

‎src/_numtype/_dtype.pyi

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import ctypes as ct
2+
from typing import Any, TypeAlias, type_check_only
3+
from typing_extensions import Protocol, TypeAliasType, TypeVar
4+
5+
import numpy as np
6+
from numpy._typing._char_codes import (
7+
_BoolCodes,
8+
_BytesCodes,
9+
_CLongDoubleCodes,
10+
_Complex64Codes,
11+
_Complex128Codes,
12+
_DT64Codes,
13+
_Float16Codes,
14+
_Float32Codes,
15+
_Float64Codes,
16+
_Int8Codes,
17+
_Int16Codes,
18+
_Int32Codes,
19+
_Int64Codes,
20+
_IntPCodes,
21+
_LongDoubleCodes,
22+
_ObjectCodes,
23+
_StrCodes,
24+
_StringCodes,
25+
_TD64Codes,
26+
_UInt8Codes,
27+
_UInt16Codes,
28+
_UInt32Codes,
29+
_UInt64Codes,
30+
_UIntPCodes,
31+
_VoidCodes,
32+
)
33+
from numpy._typing._dtype_like import _VoidDTypeLike
34+
35+
from ._just import JustBytes, JustComplex, JustFloat, JustInt, JustObject, JustStr
36+
37+
__all__ = [
38+
"ToDType",
39+
"ToDTypeBool",
40+
"ToDTypeBytes",
41+
"ToDTypeCLongDouble",
42+
"ToDTypeComplex64",
43+
"ToDTypeComplex128",
44+
"ToDTypeDateTime64",
45+
"ToDTypeFloat16",
46+
"ToDTypeFloat32",
47+
"ToDTypeFloat64",
48+
"ToDTypeInt8",
49+
"ToDTypeInt16",
50+
"ToDTypeInt32",
51+
"ToDTypeInt64",
52+
"ToDTypeLongDouble",
53+
"ToDTypeObject",
54+
"ToDTypeStr",
55+
"ToDTypeString",
56+
"ToDTypeTimeDelta64",
57+
"ToDTypeUInt8",
58+
"ToDTypeUInt16",
59+
"ToDTypeUInt32",
60+
"ToDTypeUInt64",
61+
"ToDTypeVoid",
62+
"_ToDType",
63+
"_ToDType2",
64+
]
65+
66+
_T = TypeVar("_T")
67+
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
68+
_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, covariant=True, default=Any)
69+
_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype, covariant=True, default=np.dtype)
70+
71+
@type_check_only
72+
class _HasDType(Protocol[_DTypeT_co]):
73+
@property
74+
def dtype(self) -> _DTypeT_co: ...
75+
76+
@type_check_only
77+
class _HasDTypeOf(Protocol[_ScalarT_co]):
78+
@property
79+
def dtype(self) -> np.dtype[_ScalarT_co]: ...
80+
81+
_ToDType = TypeAliasType(
82+
"_ToDType",
83+
type[_ScalarT] | np.dtype[_ScalarT] | _HasDTypeOf[_ScalarT],
84+
type_params=(_ScalarT,),
85+
)
86+
_ToDType2 = TypeAliasType(
87+
"_ToDType2",
88+
type[_ScalarT | _T] | np.dtype[_ScalarT] | _HasDTypeOf[_ScalarT],
89+
type_params=(_ScalarT, _T),
90+
)
91+
92+
_c_i16: TypeAlias = ct.c_int16 | ct.c_short
93+
_c_u16: TypeAlias = ct.c_uint16 | ct.c_ushort
94+
_c_i32: TypeAlias = ct.c_int32 | ct.c_int
95+
_c_u32: TypeAlias = ct.c_uint32 | ct.c_uint
96+
_c_i64: TypeAlias = ct.c_int64 | ct.c_longlong | ct.c_ssize_t
97+
_c_u64: TypeAlias = ct.c_uint64 | ct.c_ulonglong | ct.c_size_t | ct.c_void_p
98+
99+
ToDTypeBool: TypeAlias = _ToDType2[np.bool, ct.c_bool | bool] | _BoolCodes
100+
ToDTypeInt8: TypeAlias = _ToDType2[np.int8, ct.c_int8] | _Int8Codes
101+
ToDTypeUInt8: TypeAlias = _ToDType2[np.uint8, ct.c_uint8] | _UInt8Codes
102+
ToDTypeInt16: TypeAlias = _ToDType2[np.int16, _c_i16] | _Int16Codes
103+
ToDTypeUInt16: TypeAlias = _ToDType2[np.uint16, _c_u16] | _UInt16Codes
104+
ToDTypeInt32: TypeAlias = _ToDType2[np.int32, _c_i32] | _Int32Codes
105+
ToDTypeUInt32: TypeAlias = _ToDType2[np.uint32, _c_u32] | _UInt32Codes
106+
ToDTypeInt64: TypeAlias = _ToDType2[np.int64, _c_i64 | JustInt] | _Int64Codes | _IntPCodes
107+
ToDTypeUInt64: TypeAlias = _ToDType2[np.uint64, _c_u64] | _UInt64Codes | _UIntPCodes
108+
ToDTypeFloat16: TypeAlias = _ToDType[np.float16] | _Float16Codes
109+
ToDTypeFloat32: TypeAlias = _ToDType2[np.float32, ct.c_float] | _Float32Codes
110+
ToDTypeFloat64: TypeAlias = _ToDType2[np.float64, ct.c_double | JustFloat] | _Float64Codes
111+
ToDTypeLongDouble: TypeAlias = _ToDType2[np.longdouble, ct.c_longdouble] | _LongDoubleCodes
112+
ToDTypeComplex64: TypeAlias = _ToDType[np.complex64] | _Complex64Codes
113+
ToDTypeComplex128: TypeAlias = _ToDType2[np.complex128, JustComplex] | _Complex128Codes
114+
ToDTypeCLongDouble: TypeAlias = _ToDType[np.clongdouble] | _CLongDoubleCodes
115+
ToDTypeObject: TypeAlias = _ToDType2[np.object_, ct.py_object[Any] | JustObject] | _ObjectCodes
116+
ToDTypeBytes: TypeAlias = _ToDType2[np.bytes_, ct.c_char | JustBytes] | _BytesCodes
117+
ToDTypeStr: TypeAlias = _ToDType2[np.str_, JustStr] | _StrCodes
118+
ToDTypeVoid: TypeAlias = _ToDType2[np.void, memoryview] | _VoidDTypeLike | _VoidCodes
119+
ToDTypeDateTime64: TypeAlias = _ToDType[np.datetime64] | _DT64Codes
120+
ToDTypeTimeDelta64: TypeAlias = _ToDType[np.timedelta64] | _TD64Codes
121+
ToDTypeString: TypeAlias = np.dtypes.StringDType | _HasDType[np.dtypes.StringDType] | _StringCodes
122+
ToDType: TypeAlias = str | type | _HasDTypeOf[Any] | _VoidDTypeLike

‎src/numpy-stubs/@test/static/accept/dtype.pyi

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
import ctypes as ct
2-
from decimal import Decimal
3-
from fractions import Fraction
42
from typing_extensions import LiteralString, assert_type
53

64
import numpy as np
@@ -56,9 +54,6 @@ assert_type(np.dtype(np.number), np.dtype[np.number])
5654
# NOTE: `character` and `flexible` always fail on mypy because of some mypy bug
5755
assert_type(np.dtype(np.generic), np.dtype[np.generic])
5856

59-
assert_type(np.dtype(Decimal), np.dtype[np.object_])
60-
assert_type(np.dtype(Fraction), np.dtype[np.object_])
61-
6257
# char-codes
6358
assert_type(np.dtype("u1"), np.dtype[np.uint8])
6459
assert_type(np.dtype("int_"), np.dtype[np.intp])
@@ -79,7 +74,6 @@ assert_type(np.dtype(None), np.dtype[np.float64])
7974

8075
# Dypes of dtypes
8176
assert_type(np.dtype(np.dtype(np.float64)), np.dtype[np.float64])
82-
assert_type(np.dtype(np.dtype(np.inexact)), np.dtype[np.inexact])
8377

8478
# Parameterized dtypes
8579
assert_type(np.dtype("S8"), np.dtype)

0 commit comments

Comments
 (0)