TYP: Add support for StringDType in numpy.dtype

Author: jorenham

numpy/__init__.pyi

@@ -123,6 +123,7 @@ from numpy._typing import (
     _BytesCodes,
     _VoidCodes,
     _ObjectCodes,
+    _StringCodes,
 
     _UnsignedIntegerCodes,
     _SignedIntegerCodes,
@@ -942,6 +943,16 @@ class dtype(Generic[_DTypeScalar_co]):
     @overload
     def __new__(cls, dtype: _ObjectCodes | type[ct.py_object[Any]], align: builtins.bool = ..., copy: builtins.bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[object_]: ...
 
+    # `StringDType` requires special treatment because it has no scalar type
+    @overload
+    def __new__(
+        cls,
+        dtype: dtypes.StringDType | _StringCodes,
+        align: builtins.bool = ...,
+        copy: builtins.bool = ...,
+        metadata: dict[builtins.str, Any] = ...
+    ) -> dtypes.StringDType: ...
+
     # Combined char-codes and ctypes, analogous to the scalar-type hierarchy
     @overload
     def __new__(

numpy/_typing/__init__.py

@@ -71,6 +71,7 @@
     _BytesCodes as _BytesCodes,
     _VoidCodes as _VoidCodes,
     _ObjectCodes as _ObjectCodes,
+    _StringCodes as _StringCodes,
     _UnsignedIntegerCodes as _UnsignedIntegerCodes,
     _SignedIntegerCodes as _SignedIntegerCodes,
     _IntegerCodes as _IntegerCodes,

numpy/_typing/_char_codes.py

@@ -140,6 +140,9 @@
     "m8[as]", "|m8[as]", "=m8[as]", "<m8[as]", ">m8[as]",
 ]
 
+# NOTE: `StringDType' has no scalar type, and therefore has no name that can
+# be passed to the `dtype` constructor
+_StringCodes = Literal["T", "|T", "=T", "<T", ">T"]
 
 # NOTE: Nested literals get flattened and de-duplicated at runtime, which isn't
 # the case for a `Union` of `Literal`s.
@@ -202,4 +205,6 @@
     _DT64Codes,
     _TD64Codes,
     _ObjectCodes,
+    # TODO: add `_StringCodes` once it has a scalar type
+    # _StringCodes,
 ]

numpy/typing/tests/data/reveal/dtype.pyi

@@ -6,6 +6,7 @@ from fractions import Fraction
 from typing import Any, Literal, TypeAlias
 
 import numpy as np
+from numpy.dtypes import StringDType
 
 if sys.version_info >= (3, 11):
     from typing import assert_type
@@ -31,12 +32,13 @@ ct_floating: type[ct.c_float | ct.c_double | ct.c_longdouble]
 ct_number: type[ct.c_uint8 | ct.c_float]
 ct_generic: type[ct.c_bool | ct.c_char]
 
-cs_integer: Literal['u1', '<i2', 'L']
-cs_number: Literal['=L' ,'i', 'c16']
-cs_flex: Literal['>V', 'S']
-cs_generic: Literal['H', 'U', 'h', '|M8[Y]', '?']
+cs_integer: Literal["u1", "<i2", "L"]
+cs_number: Literal["=L" ,"i", "c16"]
+cs_flex: Literal[">V", "S"]
+cs_generic: Literal["H", "U", "h", "|M8[Y]", "?"]
 
 dt_inexact: np.dtype[np.inexact[Any]]
+dt_string: StringDType
 
 
 assert_type(np.dtype(np.float64), np.dtype[np.float64])
@@ -74,10 +76,10 @@ assert_type(np.dtype(Decimal), np.dtype[np.object_])
 assert_type(np.dtype(Fraction), np.dtype[np.object_])
 
 # char-codes
-assert_type(np.dtype('u1'), np.dtype[np.uint8])
-assert_type(np.dtype('l'), np.dtype[np.long])
-assert_type(np.dtype('longlong'), np.dtype[np.longlong])
-assert_type(np.dtype('>g'), np.dtype[np.longdouble])
+assert_type(np.dtype("u1"), np.dtype[np.uint8])
+assert_type(np.dtype("l"), np.dtype[np.long])
+assert_type(np.dtype("longlong"), np.dtype[np.longlong])
+assert_type(np.dtype(">g"), np.dtype[np.longdouble])
 assert_type(np.dtype(cs_integer), np.dtype[np.integer[Any]])
 assert_type(np.dtype(cs_number), np.dtype[np.number[Any]])
 assert_type(np.dtype(cs_flex), np.dtype[np.flexible])
@@ -104,6 +106,13 @@ assert_type(np.dtype("S8"), np.dtype[Any])
 # Void
 assert_type(np.dtype(("U", 10)), np.dtype[np.void])
 
+# StringDType
+assert_type(np.dtype(dt_string), StringDType)
+assert_type(np.dtype("T"), StringDType)
+assert_type(np.dtype("=T"), StringDType)
+assert_type(np.dtype("|T"), StringDType)
+
+
 # Methods and attributes
 assert_type(dtype_U.base, np.dtype[Any])
 assert_type(dtype_U.subdtype, None | tuple[np.dtype[Any], tuple[int, ...]])