Visitar URL original
Merge pull request #31856 from charris/backport-31846 · numpy/numpy@dc8d553 · GitHub
Skip to content

Commit dc8d553

Browse files
charrischarris
authored
Merge pull request #31856 from charris/backport-31846
BUG: fix several bugs in StringDType operations (#31846)
2 parents a007a20 + 570e25d commit dc8d553

3 files changed

Lines changed: 98 additions & 4 deletions

File tree

‎numpy/_core/src/multiarray/stringdtype/casts.cpp‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,9 @@ string_to_string_resolve_descriptors(PyObject *NPY_UNUSED(self),
147147
{
148148
if (given_descrs[1] == NULL) {
149149
loop_descrs[1] = stringdtype_finalize_descr(given_descrs[0]);
150+
if (loop_descrs[1] == NULL) {
151+
return (NPY_CASTING)-1;
152+
}
150153
}
151154
else {
152155
Py_INCREF(given_descrs[1]);
@@ -852,9 +855,13 @@ make_type2s_name(NPY_TYPES typenum) {
852855
size_t nlen = strlen(type_name);
853856

854857
const char suffix[] = "_to_StringDType";
855-
size_t slen = sizeof(prefix)/sizeof(char) - 1;
858+
size_t slen = sizeof(suffix)/sizeof(char) - 1;
856859

857860
char *buf = (char *)PyMem_RawCalloc(sizeof(char), plen + nlen + slen + 1);
861+
if (buf == NULL) {
862+
npy_gil_error(PyExc_MemoryError, "Failed allocate memory for cast");
863+
return NULL;
864+
}
858865

859866
// memcpy instead of strcpy/strncat to avoid stringop-truncation warning,
860867
// since we are not including the trailing null character
@@ -1873,15 +1880,15 @@ void_to_string(PyArrayMethod_Context *context, char *const data[],
18731880
npy_intp out_stride = strides[1];
18741881

18751882
while(N--) {
1876-
size_t out_num_bytes = utf8_buffer_size(in, max_in_size);
1883+
Py_ssize_t out_num_bytes = utf8_buffer_size(in, max_in_size);
18771884
if (out_num_bytes < 0) {
18781885
npy_gil_error(PyExc_TypeError,
18791886
"Invalid UTF-8 bytes found, cannot convert to UTF-8");
18801887
goto fail;
18811888
}
18821889
npy_static_string out_ss = {0, NULL};
18831890
if (load_new_string((npy_packed_static_string *)out,
1884-
&out_ss, out_num_bytes, allocator,
1891+
&out_ss, (size_t)out_num_bytes, allocator,
18851892
"void to string cast") == -1) {
18861893
goto fail;
18871894
}

‎numpy/_core/src/multiarray/stringdtype/dtype.c‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,7 @@ common_instance(PyArray_StringDTypeObject *dtype1, PyArray_StringDTypeObject *dt
230230
}
231231

232232
return (PyArray_StringDTypeObject *)new_stringdtype_instance(
233-
out_na_object, dtype1->coerce && dtype1->coerce);
233+
out_na_object, dtype1->coerce && dtype2->coerce);
234234
}
235235

236236
/*
@@ -266,6 +266,13 @@ as_pystring(PyObject *scalar, int coerce)
266266
Py_INCREF(scalar);
267267
return scalar;
268268
}
269+
// str subclasses (e.g. np.str_) are string data even when coercion is
270+
// disabled; convert to an exact str, preserving embedded and trailing
271+
// nulls that a round-trip through a fixed-width unicode descriptor
272+
// would strip
273+
if (PyUnicode_Check(scalar)) {
274+
return PyUnicode_FromObject(scalar);
275+
}
269276
if (coerce == 0) {
270277
PyErr_SetString(PyExc_ValueError,
271278
"StringDType only allows string data when "
@@ -635,6 +642,11 @@ stringdtype_is_known_scalar_type(PyArray_DTypeMeta *cls,
635642
{
636643
return 1;
637644
}
645+
// otherwise np.str_ discovers its fixed-width 'U' descriptor, whose
646+
// cast into StringDType strips trailing NULs setitem would preserve
647+
else if (pytype == &PyUnicodeArrType_Type) {
648+
return 1;
649+
}
638650
return 0;
639651
}
640652

@@ -654,6 +666,9 @@ stringdtype_finalize_descr(PyArray_Descr *dtype)
654666
NpyString_release_allocator(allocator);
655667
PyArray_StringDTypeObject *ret = (PyArray_StringDTypeObject *)new_stringdtype_instance(
656668
sdtype->na_object, sdtype->coerce);
669+
if (ret == NULL) {
670+
return NULL;
671+
}
657672
ret->array_owned = 1;
658673
return (PyArray_Descr *)ret;
659674
}

‎numpy/_core/tests/test_stringdtype.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,35 @@ def test_create_with_na(dtype):
144144
assert arr[1] is dtype.na_object
145145

146146

147+
def test_create_with_failing_na_comparison():
148+
# An error raised while re-creating an array-owned descriptor used to
149+
# crash with a NULL dereference in stringdtype_finalize_descr instead
150+
# of propagating.
151+
class Flaky:
152+
def __init__(self):
153+
self.calls = 0
154+
155+
def __eq__(self, other):
156+
return self is other
157+
158+
def __ne__(self, other):
159+
self.calls += 1
160+
if self.calls > 1:
161+
raise RuntimeError("boom")
162+
return False
163+
164+
def __hash__(self):
165+
return 0
166+
167+
dt = StringDType(na_object=Flaky())
168+
# the first array takes ownership of the descriptor
169+
np.array(["x"], dtype=dt)
170+
# the second array creation re-creates the descriptor, re-running the
171+
# na object classification, which fails
172+
with pytest.raises(RuntimeError, match="boom"):
173+
np.array(["y"], dtype=dt)
174+
175+
147176
@pytest.mark.parametrize("i", list(range(5)))
148177
def test_set_replace_na(i):
149178
# Test strings of various lengths can be set to NaN and then replaced.
@@ -168,6 +197,18 @@ def test_null_roundtripping():
168197
assert data[1] == arr[1]
169198

170199

200+
@pytest.mark.parametrize("coerce", [True, False])
201+
def test_np_str_trailing_nul_preserved(coerce):
202+
dtype = StringDType(coerce=coerce)
203+
value = np.str_("q\x00")
204+
arr = np.empty(1, dtype=dtype)
205+
arr[0] = value
206+
assert arr[0] == "q\x00"
207+
arr[0:1] = [value]
208+
assert arr[0] == "q\x00"
209+
assert np.array([value], dtype=dtype)[0] == "q\x00"
210+
211+
171212
@pytest.mark.parametrize(
172213
"op, pyop",
173214
[
@@ -311,6 +352,17 @@ def test_self_casts(dtype, dtype2, strings):
311352
assert_array_equal(arr[:-1], newarr[:-1])
312353

313354

355+
def test_cast_method_names():
356+
get_castingimpl = np._core._multiarray_umath._get_castingimpl
357+
string_DT = type(StringDType())
358+
for other_DT, name in [(np.dtypes.BoolDType, "bool"),
359+
(np.dtypes.Float64DType, "double")]:
360+
to_string = get_castingimpl(other_DT, string_DT)
361+
assert f"cast_{name}_to_StringDType" in repr(to_string)
362+
from_string = get_castingimpl(string_DT, other_DT)
363+
assert f"cast_StringDType_to_{name}" in repr(from_string)
364+
365+
314366
@pytest.mark.parametrize(
315367
("strings"),
316368
[
@@ -911,6 +963,15 @@ def test_string_to_bytes_invalid_ascii_error():
911963
)
912964

913965

966+
def test_void_to_string_invalid_utf8_error():
967+
# invalid UTF-8 in a void buffer used to surface as a confusing
968+
# MemoryError because the error return of utf8_buffer_size was
969+
# stored in an unsigned variable
970+
varr = np.array([b"\xff\xff\xff\xff"], dtype="V4")
971+
with pytest.raises(TypeError, match="Invalid UTF-8"):
972+
varr.astype(StringDType())
973+
974+
914975
@pytest.mark.parametrize(
915976
"typename",
916977
[
@@ -1786,6 +1847,17 @@ def test_unset_na_coercion():
17861847
arr == op
17871848

17881849

1850+
def test_coerce_promotion_commutative():
1851+
# promoting a coerce=False instance with a coerce=True instance
1852+
# used to depend on the argument order
1853+
dt = StringDType()
1854+
dt_no_coerce = StringDType(coerce=False)
1855+
assert np.promote_types(dt, dt_no_coerce) == dt_no_coerce
1856+
assert np.promote_types(dt_no_coerce, dt) == dt_no_coerce
1857+
assert np.promote_types(dt, dt) == dt
1858+
assert np.promote_types(dt_no_coerce, dt_no_coerce) == dt_no_coerce
1859+
1860+
17891861
def test_repeat(string_array):
17901862
res = string_array.repeat(1000)
17911863
# Create an empty array with expanded dimension, and fill it. Then,

0 commit comments

Comments
 (0)