Repository: numpy/numpy
Pull request: BUG: support nan-like null strings in [l,r]strip
Fixes #26390.

It turns out that `strip` has been missing from the stringdtype tests for unary and binary ufuncs and so the tests for null support were never triggered. I fixed that by adding all the strip-like functions to the tests.

Note that this adds another case where we have a `goto next_step` and an extra block to avoid compliation errors about jumping over variable declarations. I'd prefer to go through all of these and rewrite the loops to use for loops instead of `while (N--)`, this just keeps things consistent with the rest of the file.

Problem statement:
BUG: numpy.strings.lstrip/strip/rstrip fails with nulls
### Describe the issue:

When encountering a null value in numpy 2.0 StringDType'd array, we get

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/thomasli/opt/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numpy/_core/strings.py", line 881, in rstrip
    return _rstrip_whitespace(a)
ValueError: Cannot strip null values that are not strings
```

It is expected that the nulls are propagated.

### Reproduce the code example:

```python
import numpy as np
a = np.array(["a", "b", np.nan], dtype=np.dtypes.StringDType(na_object=np.nan))
np.strings.rstrip(a)
np.strings.lstrip(a)
np.strings.strip(a)
```


### Error message:

```shell
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/thomasli/opt/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numpy/_core/strings.py", line 881, in rstrip
    return _rstrip_whitespace(a)
ValueError: Cannot strip null values that are not strings

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/thomasli/opt/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numpy/_core/strings.py", line 926, in strip
    return _strip_whitespace(a)
ValueError: Cannot strip null values that are not strings

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/thomasli/opt/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numpy/_core/strings.py", line 840, in lstrip
    return _lstrip_whitespace(a)
ValueError: Cannot strip null values that are not strings
```


### Python and NumPy Versions:

2.1.0.dev0+git20240424.f8392ce

### Runtime Environment:

_No response_

### Context for the issue:

_No response_

Hints from issue discussion:
Updated the script to use `np.nan`, which should work. Thanks for the report.
Opened https://github.com/numpy/numpy/pull/26392 with a fix.

Commit message: BUG: support nan-like null strings in [l,r]strip

Patch to review:
diff --git a/numpy/_core/src/umath/stringdtype_ufuncs.cpp b/numpy/_core/src/umath/stringdtype_ufuncs.cpp
index 052c4381a4b5..8550605cc732 100644
--- a/numpy/_core/src/umath/stringdtype_ufuncs.cpp
+++ b/numpy/_core/src/umath/stringdtype_ufuncs.cpp
@@ -1046,6 +1046,7 @@ string_lrstrip_chars_strided_loop(
     PyArray_StringDTypeObject *s1descr = (PyArray_StringDTypeObject *)context->descriptors[0];
     int has_null = s1descr->na_object != NULL;
     int has_string_na = s1descr->has_string_na;
+    int has_nan_na = s1descr->has_nan_na;
 
     const npy_static_string *default_string = &s1descr->default_string;
     npy_intp N = dimensions[0];
@@ -1072,28 +1073,46 @@ string_lrstrip_chars_strided_loop(
                     s2 = *default_string;
                 }
             }
+            else if (has_nan_na) {
+                if (s2_isnull) {
+                    npy_gil_error(PyExc_ValueError,
+                                  "Cannot use a null string that is not a "
+                                  "string as the %s delimiter", ufunc_name);
+                }
+                if (s1_isnull) {
+                    if (NpyString_pack_null(oallocator, ops) < 0) {
+                        npy_gil_error(PyExc_MemoryError,
+                                      "Failed to deallocate string in %s",
+                                      ufunc_name);
+                        goto fail;
+                    }
+                    goto next_step;
+                }
+            }
             else {
                 npy_gil_error(PyExc_ValueError,
-                              "Cannot strip null values that are not strings");
+                              "Can only strip null values that are strings "
+                              "or NaN-like values");
                 goto fail;
             }
         }
+        {
+            char *new_buf = (char *)PyMem_RawCalloc(s1.size, 1);
+            Buffer<ENCODING::UTF8> buf1((char *)s1.buf, s1.size);
+            Buffer<ENCODING::UTF8> buf2((char *)s2.buf, s2.size);
+            Buffer<ENCODING::UTF8> outbuf(new_buf, s1.size);
+            size_t new_buf_size = string_lrstrip_chars
+                    (buf1, buf2, outbuf, striptype);
 
+            if (NpyString_pack(oallocator, ops, new_buf, new_buf_size) < 0) {
+                npy_gil_error(PyExc_MemoryError, "Failed to pack string in %s",
+                              ufunc_name);
+                goto fail;
+            }
 
-        char *new_buf = (char *)PyMem_RawCalloc(s1.size, 1);
-        Buffer<ENCODING::UTF8> buf1((char *)s1.buf, s1.size);
-        Buffer<ENCODING::UTF8> buf2((char *)s2.buf, s2.size);
-        Buffer<ENCODING::UTF8> outbuf(new_buf, s1.size);
-        size_t new_buf_size = string_lrstrip_chars
-                (buf1, buf2, outbuf, striptype);
-
-        if (NpyString_pack(oallocator, ops, new_buf, new_buf_size) < 0) {
-            npy_gil_error(PyExc_MemoryError, "Failed to pack string in %s",
-                          ufunc_name);
-            goto fail;
+            PyMem_RawFree(new_buf);
         }
-
-        PyMem_RawFree(new_buf);
+      next_step:
 
         in1 += strides[0];
         in2 += strides[1];
@@ -1150,8 +1169,9 @@ string_lrstrip_whitespace_strided_loop(
     const char *ufunc_name = ((PyUFuncObject *)context->caller)->name;
     STRIPTYPE striptype = *(STRIPTYPE *)context->method->static_data;
     PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)context->descriptors[0];
-    int has_string_na = descr->has_string_na;
     int has_null = descr->na_object != NULL;
+    int has_string_na = descr->has_string_na;
+    int has_nan_na = descr->has_nan_na;
     const npy_static_string *default_string = &descr->default_string;
 
     npy_string_allocator *allocators[2] = {};
@@ -1169,6 +1189,7 @@ string_lrstrip_whitespace_strided_loop(
         npy_static_string s = {0, NULL};
         int s_isnull = NpyString_load(allocator, ps, &s);
 
+
         if (s_isnull == -1) {
             npy_gil_error(PyExc_MemoryError, "Failed to load string in %s",
                           ufunc_name);
@@ -1181,26 +1202,39 @@ string_lrstrip_whitespace_strided_loop(
             if (has_string_na || !has_null) {
                 s = *default_string;
             }
+            else if (has_nan_na) {
+                if (NpyString_pack_null(oallocator, ops) < 0) {
+                    npy_gil_error(PyExc_MemoryError,
+                                  "Failed to deallocate string in %s",
+                                  ufunc_name);
+                    goto fail;
+                }
+                goto next_step;
+            }
             else {
                 npy_gil_error(PyExc_ValueError,
-                              "Cannot strip null values that are not strings");
+                              "Can only strip null values that are strings or "
+                              "NaN-like values");
                 goto fail;
             }
         }
+        {
+            char *new_buf = (char *)PyMem_RawCalloc(s.size, 1);
+            Buffer<ENCODING::UTF8> buf((char *)s.buf, s.size);
+            Buffer<ENCODING::UTF8> outbuf(new_buf, s.size);
+            size_t new_buf_size = string_lrstrip_whitespace(
+                    buf, outbuf, striptype);
 
-        char *new_buf = (char *)PyMem_RawCalloc(s.size, 1);
-        Buffer<ENCODING::UTF8> buf((char *)s.buf, s.size);
-        Buffer<ENCODING::UTF8> outbuf(new_buf, s.size);
-        size_t new_buf_size = string_lrstrip_whitespace(
-                buf, outbuf, striptype);
+            if (NpyString_pack(oallocator, ops, new_buf, new_buf_size) < 0) {
+                npy_gil_error(PyExc_MemoryError, "Failed to pack string in %s",
+                              ufunc_name);
+                goto fail;
+            }
 
-        if (NpyString_pack(oallocator, ops, new_buf, new_buf_size) < 0) {
-            npy_gil_error(PyExc_MemoryError, "Failed to pack string in %s",
-                          ufunc_name);
-            goto fail;
+            PyMem_RawFree(new_buf);
         }
 
-        PyMem_RawFree(new_buf);
+      next_step:
 
         in += strides[0];
         out += strides[1];
diff --git a/numpy/_core/tests/test_stringdtype.py b/numpy/_core/tests/test_stringdtype.py
index dd6ac36999e6..a24e643b1e64 100644
--- a/numpy/_core/tests/test_stringdtype.py
+++ b/numpy/_core/tests/test_stringdtype.py
@@ -1080,7 +1080,13 @@ def unicode_array():
     "capitalize",
     "expandtabs",
     "lower",
-    "splitlines" "swapcase" "title" "upper",
+    "lstrip",
+    "rstrip",
+    "splitlines",
+    "strip",
+    "swapcase",
+    "title",
+    "upper",
 ]
 
 BOOL_OUTPUT_FUNCTIONS = [
@@ -1107,7 +1113,10 @@ def unicode_array():
     "istitle",
     "isupper",
     "lower",
+    "lstrip",
+    "rstrip",
     "splitlines",
+    "strip",
     "swapcase",
     "title",
     "upper",
@@ -1129,10 +1138,20 @@ def unicode_array():
     "upper",
 ]
 
+ONLY_IN_NP_CHAR = [
+    "join",
+    "split",
+    "rsplit",
+    "splitlines"
+]
+
 
 @pytest.mark.parametrize("function_name", UNARY_FUNCTIONS)
 def test_unary(string_array, unicode_array, function_name):
-    func = getattr(np.char, function_name)
+    if function_name in ONLY_IN_NP_CHAR:
+        func = getattr(np.char, function_name)
+    else:
+        func = getattr(np.strings, function_name)
     dtype = string_array.dtype
     sres = func(string_array)
     ures = func(unicode_array)
@@ -1173,6 +1192,10 @@ def test_unary(string_array, unicode_array, function_name):
             with pytest.raises(ValueError):
                 func(na_arr)
         return
+    if not (is_nan or is_str):
+        with pytest.raises(ValueError):
+            func(na_arr)
+        return
     res = func(na_arr)
     if is_nan and function_name in NAN_PRESERVING_FUNCTIONS:
         assert res[0] is dtype.na_object
@@ -1197,13 +1220,17 @@ def test_unary(string_array, unicode_array, function_name):
     ("index", (None, "e")),
     ("join", ("-", None)),
     ("ljust", (None, 12)),
+    ("lstrip", (None, "A")),
     ("partition", (None, "A")),
     ("replace", (None, "A", "B")),
     ("rfind", (None, "A")),
     ("rindex", (None, "e")),
     ("rjust", (None, 12)),
+    ("rsplit", (None, "A")),
+    ("rstrip", (None, "A")),
     ("rpartition", (None, "A")),
     ("split", (None, "A")),
+    ("strip", (None, "A")),
     ("startswith", (None, "A")),
     ("zfill", (None, 12)),
 ]
@@ -1260,10 +1287,13 @@ def call_func(func, args, array, sanitize=True):
 
 @pytest.mark.parametrize("function_name, args", BINARY_FUNCTIONS)
 def test_binary(string_array, unicode_array, function_name, args):
-    func = getattr(np.char, function_name)
+    if function_name in ONLY_IN_NP_CHAR:
+        func = getattr(np.char, function_name)
+    else:
+        func = getattr(np.strings, function_name)
     sres = call_func(func, args, string_array)
     ures = call_func(func, args, unicode_array, sanitize=False)
-    if sres.dtype == StringDType():
+    if not isinstance(sres, tuple) and sres.dtype == StringDType():
         ures = ures.astype(StringDType())
     assert_array_equal(sres, ures)
 
@@ -1462,7 +1492,8 @@ def test_setup(self):
         view = self.get_view(self.a)
         sizes = np.where(is_short, view['size_and_flags'] & 0xf,
                          view['size'])
-        assert_array_equal(sizes, np.strings.str_len(self.a))
+        assert_array_equal(sizes, np.strings
+                           .str_len(self.a))
         assert_array_equal(view['xsiz'][2:],
                            np.void(b'\x00' * (self.sizeofstr // 4 - 1)))
         # Check that the medium string uses only 1 byte for its length


Code review comments:
diff_hunk: @@ -1072,28 +1073,46 @@ string_lrstrip_chars_strided_loop(
                     s2 = *default_string;
                 }
             }
+            else if (has_nan_na) {
+                if (s2_isnull) {
+                    npy_gil_error(PyExc_ValueError,
+                                  "Cannot use a null string that is not a "
+                                  "string as the %s delimiter", ufunc_name);
+                }
+                if (s1_isnull) {
+                    if (NpyString_pack_null(oallocator, ops) < 0) {
+                        npy_gil_error(PyExc_MemoryError,
+                                      "Failed to deallocate string in %s",
+                                      ufunc_name);
+                        goto fail;
+                    }
+                    goto next_step;
+                }
+            }
             else {
                 npy_gil_error(PyExc_ValueError,
-                              "Cannot strip null values that are not strings");
+                              "Can only strip null values that are strings "
+                              "or NaN-like values");
                 goto fail;
             }
         }
+        {
+            char *new_buf = (char *)PyMem_RawCalloc(s1.size, 1);
+            Buffer<ENCODING::UTF8> buf1((char *)s1.buf, s1.size);
+            Buffer<ENCODING::UTF8> buf2((char *)s2.buf, s2.size);
+            Buffer<ENCODING::UTF8> outbuf(new_buf, s1.size);
+            size_t new_buf_size = string_lrstrip_chars
+                    (buf1, buf2, outbuf, striptype);
 
+            if (NpyString_pack(oallocator, ops, new_buf, new_buf_size) < 0) {
+                npy_gil_error(PyExc_MemoryError, "Failed to pack string in %s",
+                              ufunc_name);
+                goto fail;
path: numpy/_core/src/umath/stringdtype_ufuncs.cpp
text: @user1:
Technically, missing the `PyMem_RawFree(new_buf);`?  I probably said that before, but it would seem nice to avoid reallocating a small buffer every time here (maybe directly allocating the result although that isn't possible if the operation is in-place).

@author:
Yeah, this error path indicates there's memory corruption or a bug in the implementation so a memory leak is the least of your worries if you get here. That said, easy enough to add a `PyMem_RawFree`, thanks for the catch.