Skip to content

Commit 94c4fb1

Browse files
jorenhamjorenham
authored
Merge pull request #546 from numpy/_numtype.CoString_
2 parents dbfb258 + 7ae717c commit 94c4fb1

File tree

2 files changed

+82
-45
lines changed

2 files changed

+82
-45
lines changed

‎src/_numtype/__init__.pyi

Lines changed: 81 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,10 @@ class CanLenArray(Protocol[_ScalarT_co, _ShapeT_co]):
201201
class _CanStringArray(Protocol[_ShapeT_co, _NaT_co]):
202202
def __array__(self, /) -> np.ndarray[_ShapeT_co, np.dtypes.StringDType[_NaT_co]]: ...
203203

204+
@type_check_only
205+
class _CanCoStringArray(Protocol[_ShapeT_co, _NaT_co]):
206+
def __array__(self, /) -> np.ndarray[_ShapeT_co, np.dtypes.StringDType[_NaT_co] | np.dtype[np.str_]]: ...
207+
204208
###
205209
# Shape-typed sequences
206210

@@ -256,7 +260,6 @@ _ToArray2_3nd: TypeAlias = CanLenArray[_ScalarT, Shape3N] | Sequence[_ToArray2_2
256260
###
257261
# Non-overlapping scalar- and array-like aliases for all scalar types.
258262

259-
# bool
260263
_ToBool: TypeAlias = np.bool[Any]
261264
ToBool_nd = TypeAliasType("ToBool_nd", _ToArray2_nd[_ToBool, bool])
262265
ToBool_0d = TypeAliasType("ToBool_0d", _ToArray2_0d[_ToBool, bool])
@@ -270,7 +273,6 @@ ToBool_1nd = TypeAliasType("ToBool_1nd", _ToArray2_1nd[_ToBool, bool])
270273
ToBool_2nd = TypeAliasType("ToBool_2nd", _ToArray2_2nd[_ToBool, bool])
271274
ToBool_3nd = TypeAliasType("ToBool_3nd", _ToArray2_3nd[_ToBool, bool])
272275

273-
# unsigned integers
274276
ToUInt8_nd = TypeAliasType("ToUInt8_nd", _ToArray_nd[np.uint8])
275277
ToUInt8_0d = TypeAliasType("ToUInt8_0d", CanArray0D[np.uint8])
276278
ToUInt8_1d = TypeAliasType("ToUInt8_1d", _ToArray_1d[np.uint8])
@@ -343,7 +345,6 @@ ToUInteger_1nd = TypeAliasType("ToUInteger_1nd", _ToArray_1nd[np.unsignedinteger
343345
ToUInteger_2nd = TypeAliasType("ToUInteger_2nd", _ToArray_2nd[np.unsignedinteger])
344346
ToUInteger_3nd = TypeAliasType("ToUInteger_3nd", _ToArray_3nd[np.unsignedinteger])
345347

346-
# signed integers
347348
ToInt8_nd = TypeAliasType("ToInt8_nd", _ToArray_nd[np.int8])
348349
ToInt8_0d = TypeAliasType("ToInt8_0d", CanArray0D[np.int8])
349350
ToInt8_1d = TypeAliasType("ToInt8_1d", _ToArray_1d[np.int8])
@@ -428,7 +429,6 @@ ToSInteger_1nd = TypeAliasType("ToSInteger_1nd", _ToArray2_1nd[np.signedinteger,
428429
ToSInteger_2nd = TypeAliasType("ToSInteger_2nd", _ToArray2_2nd[np.signedinteger, JustInt])
429430
ToSInteger_3nd = TypeAliasType("ToSInteger_3nd", _ToArray2_3nd[np.signedinteger, JustInt])
430431

431-
# integers
432432
ToInteger_nd = TypeAliasType("ToInteger_nd", _ToArray2_nd[np.integer, JustInt])
433433
ToInteger_0d = TypeAliasType("ToInteger_0d", _ToArray2_0d[np.integer, JustInt])
434434
ToInteger_1d = TypeAliasType("ToInteger_1d", _ToArray2_1d[np.integer, JustInt])
@@ -441,7 +441,6 @@ ToInteger_1nd = TypeAliasType("ToInteger_1nd", _ToArray2_1nd[np.integer, JustInt
441441
ToInteger_2nd = TypeAliasType("ToInteger_2nd", _ToArray2_2nd[np.integer, JustInt])
442442
ToInteger_3nd = TypeAliasType("ToInteger_3nd", _ToArray2_3nd[np.integer, JustInt])
443443

444-
# real floats
445444
ToFloat16_nd = TypeAliasType("ToFloat16_nd", _ToArray_nd[np.float16])
446445
ToFloat16_0d = TypeAliasType("ToFloat16_0d", CanArray0D[np.float16])
447446
ToFloat16_1d = TypeAliasType("ToFloat16_1d", _ToArray_1d[np.float16])
@@ -502,7 +501,6 @@ ToFloating_1nd = TypeAliasType("ToFloating_1nd", _ToArray2_1nd[np.floating, Just
502501
ToFloating_2nd = TypeAliasType("ToFloating_2nd", _ToArray2_2nd[np.floating, JustFloat])
503502
ToFloating_3nd = TypeAliasType("ToFloating_3nd", _ToArray2_3nd[np.floating, JustFloat])
504503

505-
# complex floats
506504
ToComplex64_nd = TypeAliasType("ToComplex64_nd", _ToArray_nd[np.complex64])
507505
ToComplex64_0d = TypeAliasType("ToComplex64_0d", CanArray0D[np.complex64])
508506
ToComplex64_1d = TypeAliasType("ToComplex64_1d", _ToArray_1d[np.complex64])
@@ -551,7 +549,6 @@ ToComplex_1nd = TypeAliasType("ToComplex_1nd", _ToArray2_1nd[np.complexfloating,
551549
ToComplex_2nd = TypeAliasType("ToComplex_2nd", _ToArray2_2nd[np.complexfloating, JustComplex])
552550
ToComplex_3nd = TypeAliasType("ToComplex_3nd", _ToArray2_3nd[np.complexfloating, JustComplex])
553551

554-
# integers, real- and complex floats (no booleans)
555552
ToNumber_nd = TypeAliasType("ToNumber_nd", _ToArray2_nd[np.number, _PyNumber])
556553
ToNumber_0d = TypeAliasType("ToNumber_0d", _ToArray2_0d[np.number, _PyNumber])
557554
ToNumber_1d = TypeAliasType("ToNumber_1d", _ToArray2_1d[np.number, _PyNumber])
@@ -564,7 +561,6 @@ ToNumber_1nd = TypeAliasType("ToNumber_1nd", _ToArray2_1nd[np.number, _PyNumber]
564561
ToNumber_2nd = TypeAliasType("ToNumber_2nd", _ToArray2_2nd[np.number, _PyNumber])
565562
ToNumber_3nd = TypeAliasType("ToNumber_3nd", _ToArray2_3nd[np.number, _PyNumber])
566563

567-
# integers and real floats
568564
ToReal_nd = TypeAliasType("ToReal_nd", _ToArray2_nd[np.number[float], _PyReal])
569565
ToReal_0d = TypeAliasType("ToReal_0d", _ToArray2_0d[np.number[float], _PyReal])
570566
ToReal_1d = TypeAliasType("ToReal_1d", _ToArray2_1d[np.number[float], _PyReal])
@@ -577,7 +573,6 @@ ToReal_1nd = TypeAliasType("ToReal_1nd", _ToArray2_1nd[np.number[float], _PyReal
577573
ToReal_2nd = TypeAliasType("ToReal_2nd", _ToArray2_2nd[np.number[float], _PyReal])
578574
ToReal_3nd = TypeAliasType("ToReal_3nd", _ToArray2_3nd[np.number[float], _PyReal])
579575

580-
# real- and complex floats
581576
ToInexact_nd = TypeAliasType("ToInexact_nd", _ToArray2_nd[np.inexact, _PyInexact])
582577
ToInexact_0d = TypeAliasType("ToInexact_0d", _ToArray2_0d[np.inexact, _PyInexact])
583578
ToInexact_1d = TypeAliasType("ToInexact_1d", _ToArray2_1d[np.inexact, _PyInexact])
@@ -590,7 +585,6 @@ ToInexact_1nd = TypeAliasType("ToInexact_1nd", _ToArray2_1nd[np.inexact, _PyInex
590585
ToInexact_2nd = TypeAliasType("ToInexact_2nd", _ToArray2_2nd[np.inexact, _PyInexact])
591586
ToInexact_3nd = TypeAliasType("ToInexact_3nd", _ToArray2_3nd[np.inexact, _PyInexact])
592587

593-
# temporal
594588
_ToTimeDelta: TypeAlias = np.timedelta64[Any]
595589
ToTimeDelta_nd = TypeAliasType("ToTimeDelta_nd", _ToArray_nd[_ToTimeDelta])
596590
ToTimeDelta_0d = TypeAliasType("ToTimeDelta_0d", CanArray0D[_ToTimeDelta])
@@ -617,7 +611,6 @@ ToDateTime_1nd = TypeAliasType("ToDateTime_1nd", _ToArray_1nd[_ToDateTime])
617611
ToDateTime_2nd = TypeAliasType("ToDateTime_2nd", _ToArray_2nd[_ToDateTime])
618612
ToDateTime_3nd = TypeAliasType("ToDateTime_3nd", _ToArray_3nd[_ToDateTime])
619613

620-
# fixed strings
621614
_ToBytes: TypeAlias = np.character[bytes]
622615
ToBytes_nd = TypeAliasType("ToBytes_nd", _ToArray2_nd[_ToBytes, JustBytes])
623616
ToBytes_0d = TypeAliasType("ToBytes_0d", _ToArray2_0d[_ToBytes, JustBytes])
@@ -657,22 +650,13 @@ ToCharacter_1nd = TypeAliasType("ToCharacter_1nd", _ToArray2_1nd[_ToCharacter, _
657650
ToCharacter_2nd = TypeAliasType("ToCharacter_2nd", _ToArray2_2nd[_ToCharacter, _PyCharacter])
658651
ToCharacter_3nd = TypeAliasType("ToCharacter_3nd", _ToArray2_3nd[_ToCharacter, _PyCharacter])
659652

660-
# python object
661-
ToObject_nd = TypeAliasType("ToObject_nd", _ToArray2_nd[np.object_, _PyObject])
662-
ToObject_0d = TypeAliasType("ToObject_0d", _ToArray2_0d[np.object_, _PyObject])
663-
ToObject_1d = TypeAliasType("ToObject_1d", _ToArray2_1d[np.object_, _PyObject])
664-
ToObject_2d = TypeAliasType("ToObject_2d", _ToArray2_2d[np.object_, _PyObject])
665-
ToObject_3d = TypeAliasType("ToObject_3d", _ToArray2_3d[np.object_, _PyObject])
666-
ToObject_1ds = TypeAliasType("ToObject_1ds", _ToArray2_1ds[np.object_, _PyObject])
667-
ToObject_2ds = TypeAliasType("ToObject_2ds", _ToArray2_2ds[np.object_, _PyObject])
668-
ToObject_3ds = TypeAliasType("ToObject_3ds", _ToArray2_3ds[np.object_, _PyObject])
669-
ToObject_1nd = TypeAliasType("ToObject_1nd", _ToArray2_1nd[np.object_, _PyObject])
670-
ToObject_2nd = TypeAliasType("ToObject_2nd", _ToArray2_2nd[np.object_, _PyObject])
671-
ToObject_3nd = TypeAliasType("ToObject_3nd", _ToArray2_3nd[np.object_, _PyObject])
672-
673-
# StringDType
674653
ToString_nd = TypeAliasType("ToString_nd", SequenceND[_CanStringArray[Shape, _NaT0]], type_params=(_NaT0,))
675-
ToString_1ds = TypeAliasType("ToString_1ds", _CanStringArray[Shape1, _NaT0], type_params=(_NaT0,))
654+
ToString_0d = TypeAliasType("ToString_0d", _CanStringArray[Shape0, _NaT0], type_params=(_NaT0,))
655+
ToString_1ds = TypeAliasType(
656+
"ToString_1ds",
657+
_CanStringArray[Shape1, _NaT0] | Sequence[ToString_0d[_NaT0]],
658+
type_params=(_NaT0,),
659+
)
676660
ToString_2ds = TypeAliasType(
677661
"ToString_2ds",
678662
_CanStringArray[Shape2, _NaT0] | Sequence[ToString_1ds[_NaT0]],
@@ -683,11 +667,34 @@ ToString_3ds = TypeAliasType(
683667
_CanStringArray[Shape3, _NaT0] | Sequence[ToString_2ds[_NaT0]],
684668
type_params=(_NaT0,),
685669
)
686-
ToString_1nd = TypeAliasType("ToString_1nd", _CanStringArray[Shape1N, _NaT0], type_params=(_NaT0,))
687-
ToString_2nd = TypeAliasType("ToString_2nd", _CanStringArray[Shape2N, _NaT0], type_params=(_NaT0,))
688-
ToString_3nd = TypeAliasType("ToString_3nd", _CanStringArray[Shape3N, _NaT0], type_params=(_NaT0,))
670+
ToString_1nd = TypeAliasType(
671+
"ToString_1nd",
672+
_CanStringArray[Shape1N, _NaT0] | Sequence[_CanStringArray[Shape0N, _NaT0]],
673+
type_params=(_NaT0,),
674+
)
675+
ToString_2nd = TypeAliasType(
676+
"ToString_2nd",
677+
_CanStringArray[Shape2N, _NaT0] | Sequence[ToString_1nd[_NaT0]],
678+
type_params=(_NaT0,),
679+
)
680+
ToString_3nd = TypeAliasType(
681+
"ToString_3nd",
682+
_CanStringArray[Shape3N, _NaT0] | Sequence[ToString_2nd[_NaT0]],
683+
type_params=(_NaT0,),
684+
)
685+
686+
ToObject_nd = TypeAliasType("ToObject_nd", _ToArray2_nd[np.object_, _PyObject])
687+
ToObject_0d = TypeAliasType("ToObject_0d", _ToArray2_0d[np.object_, _PyObject])
688+
ToObject_1d = TypeAliasType("ToObject_1d", _ToArray2_1d[np.object_, _PyObject])
689+
ToObject_2d = TypeAliasType("ToObject_2d", _ToArray2_2d[np.object_, _PyObject])
690+
ToObject_3d = TypeAliasType("ToObject_3d", _ToArray2_3d[np.object_, _PyObject])
691+
ToObject_1ds = TypeAliasType("ToObject_1ds", _ToArray2_1ds[np.object_, _PyObject])
692+
ToObject_2ds = TypeAliasType("ToObject_2ds", _ToArray2_2ds[np.object_, _PyObject])
693+
ToObject_3ds = TypeAliasType("ToObject_3ds", _ToArray2_3ds[np.object_, _PyObject])
694+
ToObject_1nd = TypeAliasType("ToObject_1nd", _ToArray2_1nd[np.object_, _PyObject])
695+
ToObject_2nd = TypeAliasType("ToObject_2nd", _ToArray2_2nd[np.object_, _PyObject])
696+
ToObject_3nd = TypeAliasType("ToObject_3nd", _ToArray2_3nd[np.object_, _PyObject])
689697

690-
# any scalar
691698
ToGeneric_nd = TypeAliasType("ToGeneric_nd", _ToArray2_nd[np.generic, _PyScalar])
692699
ToGeneric_0d = TypeAliasType("ToGeneric_0d", _ToArray2_0d[np.generic, _PyScalar])
693700
ToGeneric_1d = TypeAliasType("ToGeneric_1d", _ToArray2_1d[np.generic, _PyScalar])
@@ -700,7 +707,9 @@ ToGeneric_1nd = TypeAliasType("ToGeneric_1nd", _ToArray2_1nd[np.generic, _PyScal
700707
ToGeneric_2nd = TypeAliasType("ToGeneric_2nd", _ToArray2_2nd[np.generic, _PyScalar])
701708
ToGeneric_3nd = TypeAliasType("ToGeneric_3nd", _ToArray2_3nd[np.generic, _PyScalar])
702709

703-
# unsigned integers
710+
###
711+
# *Co*ercible (overlapping) scalar- and array-like aliases.
712+
704713
CoUInt8_nd = TypeAliasType("CoUInt8_nd", _ToArray2_nd[co_uint8, bool])
705714
CoUInt8_0d = TypeAliasType("CoUInt8_0d", _ToArray2_0d[co_uint8, bool])
706715
CoUInt8_1d = TypeAliasType("CoUInt8_1d", _ToArray2_1d[co_uint8, bool])
@@ -761,7 +770,6 @@ CoUInt64_1nd = TypeAliasType("CoUInt64_1nd", _ToArray2_1nd[co_uint64, bool])
761770
CoUInt64_2nd = TypeAliasType("CoUInt64_2nd", _ToArray2_2nd[co_uint64, bool])
762771
CoUInt64_3nd = TypeAliasType("CoUInt64_3nd", _ToArray2_3nd[co_uint64, bool])
763772

764-
# signed integers
765773
CoInt8_nd = TypeAliasType("CoInt8_nd", _ToArray2_nd[co_int8, bool])
766774
CoInt8_0d = TypeAliasType("CoInt8_0d", _ToArray2_0d[co_int8, bool])
767775
CoInt8_1d = TypeAliasType("CoInt8_1d", _ToArray2_1d[co_int8, bool])
@@ -834,8 +842,6 @@ CoInteger_1nd = TypeAliasType("CoInteger_1nd", _ToArray2_1nd[co_integer, int])
834842
CoInteger_2nd = TypeAliasType("CoInteger_2nd", _ToArray2_2nd[co_integer, int])
835843
CoInteger_3nd = TypeAliasType("CoInteger_3nd", _ToArray2_3nd[co_integer, int])
836844

837-
# real floats
838-
839845
CoFloat16_nd = TypeAliasType("CoFloat16_nd", _ToArray2_nd[co_float16, bool])
840846
CoFloat16_0d = TypeAliasType("CoFloat16_0d", _ToArray2_0d[co_float16, bool])
841847
CoFloat16_1d = TypeAliasType("CoFloat16_1d", _ToArray2_1d[co_float16, bool])
@@ -884,7 +890,6 @@ CoFloating_1nd = TypeAliasType("CoFloating_1nd", _ToArray2_1nd[co_float, float])
884890
CoFloating_2nd = TypeAliasType("CoFloating_2nd", _ToArray2_2nd[co_float, float])
885891
CoFloating_3nd = TypeAliasType("CoFloating_3nd", _ToArray2_3nd[co_float, float])
886892

887-
# complex floats
888893
CoComplex64_nd = TypeAliasType("CoComplex64_nd", _ToArray2_nd[co_complex64, bool])
889894
CoComplex64_0d = TypeAliasType("CoComplex64_0d", _ToArray2_0d[co_complex64, bool])
890895
CoComplex64_1d = TypeAliasType("CoComplex64_1d", _ToArray2_1d[co_complex64, bool])
@@ -921,7 +926,6 @@ CoComplex_1nd = TypeAliasType("CoComplex_1nd", _ToArray2_1nd[co_complex, complex
921926
CoComplex_2nd = TypeAliasType("CoComplex_2nd", _ToArray2_2nd[co_complex, complex])
922927
CoComplex_3nd = TypeAliasType("CoComplex_3nd", _ToArray2_3nd[co_complex, complex])
923928

924-
# temporal
925929
CoTimeDelta_nd = TypeAliasType("CoTimeDelta_nd", _ToArray2_nd[co_timedelta, int])
926930
CoTimeDelta_0d = TypeAliasType("CoTimeDelta_0d", _ToArray2_0d[co_timedelta, int])
927931
CoTimeDelta_1d = TypeAliasType("CoTimeDelta_1d", _ToArray2_1d[co_timedelta, int])
@@ -946,7 +950,6 @@ CoDateTime_1nd = TypeAliasType("CoDateTime_1nd", _ToArray_1nd[co_datetime])
946950
CoDateTime_2nd = TypeAliasType("CoDateTime_2nd", _ToArray_2nd[co_datetime])
947951
CoDateTime_3nd = TypeAliasType("CoDateTime_3nd", _ToArray_3nd[co_datetime])
948952

949-
# fixed strings
950953
CoBytes_nd = TypeAliasType("CoBytes_nd", _ToArray2_nd[_ToBytes, JustBytes])
951954
CoBytes_0d = TypeAliasType("CoBytes_0d", _ToArray2_0d[_ToBytes, JustBytes])
952955
CoBytes_1d = TypeAliasType("CoBytes_1d", _ToArray2_1d[_ToBytes, JustBytes])
@@ -970,3 +973,44 @@ CoStr_3ds = TypeAliasType("CoStr_3ds", _ToArray2_3ds[_ToCharacter, _PyCharacter]
970973
CoStr_1nd = TypeAliasType("CoStr_1nd", _ToArray2_1nd[_ToCharacter, _PyCharacter])
971974
CoStr_2nd = TypeAliasType("CoStr_2nd", _ToArray2_2nd[_ToCharacter, _PyCharacter])
972975
CoStr_3nd = TypeAliasType("CoStr_3nd", _ToArray2_3nd[_ToCharacter, _PyCharacter])
976+
977+
CoString_nd = TypeAliasType(
978+
"CoString_nd",
979+
SequenceND[_CanCoStringArray[Shape, _NaT0] | JustStr],
980+
type_params=(_NaT0,),
981+
)
982+
CoString_0d = TypeAliasType(
983+
"CoString_0d",
984+
_CanCoStringArray[Shape0, _NaT0] | JustStr,
985+
type_params=(_NaT0,),
986+
)
987+
CoString_1ds = TypeAliasType(
988+
"CoString_1ds",
989+
_CanCoStringArray[Shape1, _NaT0] | Sequence[JustStr],
990+
type_params=(_NaT0,),
991+
)
992+
CoString_2ds = TypeAliasType(
993+
"CoString_2ds",
994+
_CanCoStringArray[Shape2, _NaT0] | Sequence[CoString_1ds[_NaT0]],
995+
type_params=(_NaT0,),
996+
)
997+
CoString_3ds = TypeAliasType(
998+
"CoString_3ds",
999+
_CanCoStringArray[Shape3, _NaT0] | Sequence[CoString_2ds[_NaT0]],
1000+
type_params=(_NaT0,),
1001+
)
1002+
CoString_1nd = TypeAliasType(
1003+
"CoString_1nd",
1004+
_CanCoStringArray[Shape1N, _NaT0] | Sequence[_CanCoStringArray[Shape0N, _NaT0]],
1005+
type_params=(_NaT0,),
1006+
)
1007+
CoString_2nd = TypeAliasType(
1008+
"CoString_2nd",
1009+
_CanCoStringArray[Shape2N, _NaT0] | Sequence[CoString_1nd[_NaT0]],
1010+
type_params=(_NaT0,),
1011+
)
1012+
CoString_3nd = TypeAliasType(
1013+
"CoString_3nd",
1014+
_CanCoStringArray[Shape3N, _NaT0] | Sequence[CoString_2nd[_NaT0]],
1015+
type_params=(_NaT0,),
1016+
)

‎src/numpy-stubs/_core/umath.pyi

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -188,14 +188,7 @@ _UFuncMethod: TypeAlias = L["__call__", "reduce", "reduceat", "accumulate", "out
188188
_TimeLike: TypeAlias = np.datetime64 | np.timedelta64
189189

190190
_SupportsStringLikeArray: TypeAlias = _SupportsArray[np.dtypes.StringDType | np.dtype[np.character]]
191-
_ToCharStringND: TypeAlias = (
192-
_SupportsStringLikeArray
193-
| _NestedSequence[_SupportsStringLikeArray]
194-
| list[str]
195-
| list[bytes]
196-
| _NestedSequence[list[str]]
197-
| _NestedSequence[list[bytes]]
198-
)
191+
_ToCharStringND: TypeAlias = _nt.SequenceND[_SupportsStringLikeArray | list[str] | list[bytes] | list[str | bytes]]
199192

200193
###
201194
# helper protocols

0 commit comments

Comments
 (0)