Skip to content

Commit c04f1e8

Browse files
MarcoGorellicharris
MarcoGorelli
authored and
charris
committed
TYP: Fix overloads in np.char.array and np.char.asarray for unicode argument (#29377)
1 parent 20d973c commit c04f1e8

File tree

2 files changed

+66
-10
lines changed

2 files changed

+66
-10
lines changed

‎numpy/_core/defchararray.pyi‎

Lines changed: 60 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1020,66 +1020,116 @@ def startswith(
10201020
def str_len(A: UST_co) -> NDArray[int_]: ...
10211021

10221022
# Overload 1 and 2: str- or bytes-based array-likes
1023-
# overload 3: arbitrary object with unicode=False (-> bytes_)
1024-
# overload 4: arbitrary object with unicode=True (-> str_)
1023+
# overload 3 and 4: arbitrary object with unicode=False (-> bytes_)
1024+
# overload 5 and 6: arbitrary object with unicode=True (-> str_)
1025+
# overload 7: arbitrary object with unicode=None (default) (-> str_ | bytes_)
10251026
@overload
10261027
def array(
10271028
obj: U_co,
10281029
itemsize: int | None = ...,
10291030
copy: bool = ...,
1030-
unicode: L[False] = ...,
1031+
unicode: L[True] | None = ...,
10311032
order: _OrderKACF = ...,
10321033
) -> _CharArray[str_]: ...
10331034
@overload
10341035
def array(
10351036
obj: S_co,
10361037
itemsize: int | None = ...,
10371038
copy: bool = ...,
1038-
unicode: L[False] = ...,
1039+
unicode: L[False] | None = ...,
1040+
order: _OrderKACF = ...,
1041+
) -> _CharArray[bytes_]: ...
1042+
@overload
1043+
def array(
1044+
obj: object,
1045+
itemsize: int | None,
1046+
copy: bool,
1047+
unicode: L[False],
10391048
order: _OrderKACF = ...,
10401049
) -> _CharArray[bytes_]: ...
10411050
@overload
10421051
def array(
10431052
obj: object,
10441053
itemsize: int | None = ...,
10451054
copy: bool = ...,
1046-
unicode: L[False] = ...,
1055+
*,
1056+
unicode: L[False],
10471057
order: _OrderKACF = ...,
10481058
) -> _CharArray[bytes_]: ...
10491059
@overload
1060+
def array(
1061+
obj: object,
1062+
itemsize: int | None,
1063+
copy: bool,
1064+
unicode: L[True],
1065+
order: _OrderKACF = ...,
1066+
) -> _CharArray[str_]: ...
1067+
@overload
10501068
def array(
10511069
obj: object,
10521070
itemsize: int | None = ...,
10531071
copy: bool = ...,
1054-
unicode: L[True] = ...,
1072+
*,
1073+
unicode: L[True],
10551074
order: _OrderKACF = ...,
10561075
) -> _CharArray[str_]: ...
1076+
@overload
1077+
def array(
1078+
obj: object,
1079+
itemsize: int | None = ...,
1080+
copy: bool = ...,
1081+
unicode: bool | None = ...,
1082+
order: _OrderKACF = ...,
1083+
) -> _CharArray[str_] | _CharArray[bytes_]: ...
10571084

10581085
@overload
10591086
def asarray(
10601087
obj: U_co,
10611088
itemsize: int | None = ...,
1062-
unicode: L[False] = ...,
1089+
unicode: L[True] | None = ...,
10631090
order: _OrderKACF = ...,
10641091
) -> _CharArray[str_]: ...
10651092
@overload
10661093
def asarray(
10671094
obj: S_co,
10681095
itemsize: int | None = ...,
1069-
unicode: L[False] = ...,
1096+
unicode: L[False] | None = ...,
1097+
order: _OrderKACF = ...,
1098+
) -> _CharArray[bytes_]: ...
1099+
@overload
1100+
def asarray(
1101+
obj: object,
1102+
itemsize: int | None,
1103+
unicode: L[False],
10701104
order: _OrderKACF = ...,
10711105
) -> _CharArray[bytes_]: ...
10721106
@overload
10731107
def asarray(
10741108
obj: object,
10751109
itemsize: int | None = ...,
1076-
unicode: L[False] = ...,
1110+
*,
1111+
unicode: L[False],
10771112
order: _OrderKACF = ...,
10781113
) -> _CharArray[bytes_]: ...
10791114
@overload
1115+
def asarray(
1116+
obj: object,
1117+
itemsize: int | None,
1118+
unicode: L[True],
1119+
order: _OrderKACF = ...,
1120+
) -> _CharArray[str_]: ...
1121+
@overload
10801122
def asarray(
10811123
obj: object,
10821124
itemsize: int | None = ...,
1083-
unicode: L[True] = ...,
1125+
*,
1126+
unicode: L[True],
10841127
order: _OrderKACF = ...,
10851128
) -> _CharArray[str_]: ...
1129+
@overload
1130+
def asarray(
1131+
obj: object,
1132+
itemsize: int | None = ...,
1133+
unicode: bool | None = ...,
1134+
order: _OrderKACF = ...,
1135+
) -> _CharArray[str_] | _CharArray[bytes_]: ...

‎numpy/typing/tests/data/reveal/char.pyi‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,10 +209,16 @@ assert_type(np.char.array("bob", copy=True), np.char.chararray[np_t._AnyShape, n
209209
assert_type(np.char.array(b"bob", itemsize=5), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
210210
assert_type(np.char.array(1, unicode=False), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
211211
assert_type(np.char.array(1, unicode=True), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])
212+
assert_type(np.char.array(1), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]] | np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
213+
assert_type(np.char.array(AR_U, unicode=False), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
214+
assert_type(np.char.array(AR_S, unicode=True), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])
212215

213216
assert_type(np.char.asarray(AR_U), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])
214217
assert_type(np.char.asarray(AR_S, order="K"), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
215218
assert_type(np.char.asarray("bob"), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])
216219
assert_type(np.char.asarray(b"bob", itemsize=5), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
217220
assert_type(np.char.asarray(1, unicode=False), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
218221
assert_type(np.char.asarray(1, unicode=True), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])
222+
assert_type(np.char.asarray(1), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]] | np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
223+
assert_type(np.char.asarray(AR_U, unicode=False), np.char.chararray[np_t._AnyShape, np.dtype[np.bytes_]])
224+
assert_type(np.char.asarray(AR_S, unicode=True), np.char.chararray[np_t._AnyShape, np.dtype[np.str_]])

0 commit comments

Comments
 (0)