BUG: ensure np.vectorize doesn't truncate fixed-width strings

ngoldbaum · ngoldbaum · commit 9796349b455b · 2024-04-12T08:43:57.000-06:00
diff --git a/numpy/_core/tests/test_regression.py b/numpy/_core/tests/test_regression.py
@@ -2614,3 +2614,9 @@ def test_logspace_base_does_not_determine_dtype(self):
                            base=np.array([10.0]))
         with pytest.raises(AssertionError, match="not almost equal"):
             assert_almost_equal(out2, expected)
+
+    def test_vectorize_fixed_width_string(self):
+        arr = np.array(["SOme wOrd Ǆ ß ᾛ ΣΣ ﬃ⁵Å Ç Ⅰ"]).astype(np.str_)
+        f = str.casefold
+        res = np.vectorize(f, otypes=[arr.dtype])(arr)
+        assert res.dtype == "U30"
diff --git a/numpy/lib/_function_base_impl.py b/numpy/lib/_function_base_impl.py
@@ -2132,6 +2132,12 @@ def _create_arrays(broadcast_shape, dim_sizes, list_of_core_dims, dtypes,
     return arrays
 
 
+def _get_vectorize_dtype(dtype):
+    if dtype.char in "SU":
+        return dtype.char
+    return dtype
+
+
 @set_module('numpy')
 class vectorize:
     """
@@ -2330,7 +2336,7 @@ def __init__(self, pyfunc=np._NoValue, otypes=None, doc=None,
                 if char not in typecodes['All']:
                     raise ValueError("Invalid otype specified: %s" % (char,))
         elif iterable(otypes):
-            otypes = [_nx.dtype(x) for x in otypes]
+            otypes = [_get_vectorize_dtype(_nx.dtype(x)) for x in otypes]
         elif otypes is not None:
             raise ValueError("Invalid otype specification")
         self.otypes = otypes
