BUG: add bounds-checking to in-place string multiply

ngoldbaum · ngoldbaum · commit db879703b7ed · 2025-05-26T19:10:14.000-06:00
diff --git a/numpy/_core/src/umath/string_buffer.h b/numpy/_core/src/umath/string_buffer.h
@@ -297,6 +297,18 @@ struct Buffer {
         return num_codepoints;
     }
 
+    inline size_t
+    buffer_width()
+    {
+        switch (enc) {
+            case ENCODING::ASCII:
+            case ENCODING::UTF8:
+                return after - buf;
+            case ENCODING::UTF32:
+                return (after - buf) / sizeof(npy_ucs4);
+        }
+    }
+
     inline Buffer<enc>&
     operator+=(npy_int64 rhs)
     {
diff --git a/numpy/_core/src/umath/string_ufuncs.cpp b/numpy/_core/src/umath/string_ufuncs.cpp
@@ -175,6 +175,12 @@ string_multiply(Buffer<enc> buf1, npy_int64 reps, Buffer<enc> out)
         return;
     }
 
+    size_t width = out.buffer_width();
+    size_t pad = 0;
+    if (width < len1*reps) {
+        reps = width / len1;
+        pad = width % len1;
+    }
     if (len1 == 1) {
         out.buffer_memset(*buf1, reps);
         out.buffer_fill_with_zeros_after_index(reps);
@@ -184,6 +190,8 @@ string_multiply(Buffer<enc> buf1, npy_int64 reps, Buffer<enc> out)
             buf1.buffer_memcpy(out, len1);
             out += len1;
         }
+        buf1.buffer_memcpy(out, pad);
+        out += pad;
         out.buffer_fill_with_zeros_after_index(0);
     }
 }
@@ -752,10 +760,11 @@ string_multiply_resolve_descriptors(
     if (given_descrs[2] == NULL) {
         PyErr_SetString(
             PyExc_TypeError,
-            "The 'out' kwarg is necessary. Use numpy.strings.multiply without it.");
+            "The 'out' kwarg is necessary when using the string multiply ufunc "
+            "directly. Use numpy.strings.multiply to multiply strings without "
+            "specifying 'out'.");
         return _NPY_ERROR_OCCURRED_IN_CAST;
     }
-
     loop_descrs[0] = NPY_DT_CALL_ensure_canonical(given_descrs[0]);
     if (loop_descrs[0] == NULL) {
         return _NPY_ERROR_OCCURRED_IN_CAST;
diff --git a/numpy/_core/tests/test_strings.py b/numpy/_core/tests/test_strings.py
@@ -227,6 +227,14 @@ def test_multiply_raises(self, dt):
         with pytest.raises(MemoryError):
             np.strings.multiply(np.array("abc", dtype=dt), sys.maxsize)
 
+    def test_inplace_multiply(self, dt):
+        arr = np.array(['foo ', 'bar'], dtype=dt)
+        arr *= 2
+        if dt != "T":
+            assert_array_equal(arr, np.array(['foo ', 'barb'], dtype=dt))
+        else:
+            assert_array_equal(arr, ['foo foo ', 'barbar'])
+
     @pytest.mark.parametrize("i_dt", [np.int8, np.int16, np.int32,
                                       np.int64, np.int_])
     def test_multiply_integer_dtypes(self, i_dt, dt):
diff --git a/numpy/typing/tests/data/pass/ma.py b/numpy/typing/tests/data/pass/ma.py
@@ -16,7 +16,8 @@
 MAR_M_dt64: MaskedArray[np.datetime64] = np.ma.MaskedArray([np.datetime64(1, "D")])
 MAR_S: MaskedArray[np.bytes_] = np.ma.MaskedArray([b'foo'], dtype=np.bytes_)
 MAR_U: MaskedArray[np.str_] = np.ma.MaskedArray(['foo'], dtype=np.str_)
-MAR_T = cast(np.ma.MaskedArray[Any, np.dtypes.StringDType], np.ma.MaskedArray(["a"], "T"))
+MAR_T = cast(np.ma.MaskedArray[Any, np.dtypes.StringDType],
+             np.ma.MaskedArray(["a"], dtype="T"))
 
 AR_b: npt.NDArray[np.bool] = np.array([True, False, True])
 
