BUG: fix more issues with string ufunc promotion

Author: ngoldbaum

numpy/_core/strings.py

@@ -669,20 +669,29 @@ def center(a, width, fillchar=' '):
     array(['a1b2', '1b2a', 'b2a1', '2a1b'], dtype='<U4')
 
     """
+    width = np.asanyarray(width)
+    if not np.issubdtype(width.dtype, np.integer):
+        raise TypeError(f"unsupported type {width.dtype} for operand 'width'")
+
     a = np.asanyarray(a)
-    fillchar = np.asanyarray(fillchar, dtype=a.dtype)
+    fillchar = np.asanyarray(fillchar)
+
+    try_out_dt = np.result_type(a, fillchar)
+    if try_out_dt.char == "T":
+        a = a.astype(try_out_dt, copy=False)
+        fillchar = fillchar.astype(try_out_dt, copy=False)
+        out = None
+    else:
+        fillchar = fillchar.astype(a.dtype, copy=False)
+        width = np.maximum(str_len(a), width)
+        out_dtype = f"{a.dtype.char}{width.max()}"
+        shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
+        out = np.empty_like(a, shape=shape, dtype=out_dtype)
 
     if np.any(str_len(fillchar) != 1):
         raise TypeError(
             "The fill character must be exactly one character long")
 
-    if a.dtype.char == "T":
-        return _center(a, width, fillchar)
-
-    width = np.maximum(str_len(a), width)
-    out_dtype = f"{a.dtype.char}{width.max()}"
-    shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
-    out = np.empty_like(a, shape=shape, dtype=out_dtype)
     return _center(a, width, fillchar, out=out)
 
 
@@ -726,20 +735,29 @@ def ljust(a, width, fillchar=' '):
     array(['aAaAaA   ', '  aA     ', 'abBABba  '], dtype='<U9')
 
     """
+    width = np.asanyarray(width)
+    if not np.issubdtype(width.dtype, np.integer):
+        raise TypeError(f"unsupported type {width.dtype} for operand 'width'")
+
     a = np.asanyarray(a)
-    fillchar = np.asanyarray(fillchar, dtype=a.dtype)
+    fillchar = np.asanyarray(fillchar)
+
+    try_out_dt = np.result_type(a, fillchar)
+    if try_out_dt.char == "T":
+        a = a.astype(try_out_dt, copy=False)
+        fillchar = fillchar.astype(try_out_dt, copy=False)
+        out = None
+    else:
+        fillchar = fillchar.astype(a.dtype, copy=False)
+        width = np.maximum(str_len(a), width)
+        shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
+        out_dtype = f"{a.dtype.char}{width.max()}"
+        out = np.empty_like(a, shape=shape, dtype=out_dtype)
 
     if np.any(str_len(fillchar) != 1):
         raise TypeError(
             "The fill character must be exactly one character long")
 
-    if a.dtype.char == "T":
-        return _ljust(a, width, fillchar)
-
-    width = np.maximum(str_len(a), width)
-    shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
-    out_dtype = f"{a.dtype.char}{width.max()}"
-    out = np.empty_like(a, shape=shape, dtype=out_dtype)
     return _ljust(a, width, fillchar, out=out)
 
 
@@ -783,20 +801,29 @@ def rjust(a, width, fillchar=' '):
     array(['   aAaAaA', '     aA  ', '  abBABba'], dtype='<U9')
 
     """
+    width = np.asanyarray(width)
+    if not np.issubdtype(width.dtype, np.integer):
+        raise TypeError(f"unsupported type {width.dtype} for operand 'width'")
+
     a = np.asanyarray(a)
-    fillchar = np.asanyarray(fillchar, dtype=a.dtype)
+    fillchar = np.asanyarray(fillchar)
+
+    try_out_dt = np.result_type(a, fillchar)
+    if try_out_dt.char == "T":
+        a = a.astype(try_out_dt, copy=False)
+        fillchar = fillchar.astype(try_out_dt, copy=False)
+        out = None
+    else:
+        fillchar = fillchar.astype(a.dtype, copy=False)
+        width = np.maximum(str_len(a), width)
+        shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
+        out_dtype = f"{a.dtype.char}{width.max()}"
+        out = np.empty_like(a, shape=shape, dtype=out_dtype)
 
     if np.any(str_len(fillchar) != 1):
         raise TypeError(
             "The fill character must be exactly one character long")
 
-    if a.dtype.char == "T":
-        return _rjust(a, width, fillchar)
-
-    width = np.maximum(str_len(a), width)
-    shape = np.broadcast_shapes(a.shape, width.shape, fillchar.shape)
-    out_dtype = f"{a.dtype.char}{width.max()}"
-    out = np.empty_like(a, shape=shape, dtype=out_dtype)
     return _rjust(a, width, fillchar, out=out)
 
 
@@ -830,6 +857,10 @@ def zfill(a, width):
     array(['001', '-01', '+01'], dtype='<U3')
 
     """
+    width = np.asanyarray(width)
+    if not np.issubdtype(width.dtype, np.integer):
+        raise TypeError(f"unsupported type {width.dtype} for operand 'width'")
+
     a = np.asanyarray(a)
 
     if a.dtype.char == "T":
@@ -1205,22 +1236,33 @@ def replace(a, old, new, count=-1):
     array(['The dwash was fresh', 'Thwas was it'], dtype='<U19')
 
     """
-    arr = np.asanyarray(a)
-    a_dt = arr.dtype
-    old = np.asanyarray(old, dtype=getattr(old, 'dtype', a_dt))
-    new = np.asanyarray(new, dtype=getattr(new, 'dtype', a_dt))
     count = np.asanyarray(count)
+    if not np.issubdtype(count.dtype, np.integer):
+        raise TypeError(f"unsupported type {count.dtype} for operand 'count'")
 
-    if arr.dtype.char == "T":
-        return _replace(arr, old, new, count)
-
-    max_int64 = np.iinfo(np.int64).max
-    counts = _count_ufunc(arr, old, 0, max_int64)
-    counts = np.where(count < 0, counts, np.minimum(counts, count))
-
-    buffersizes = str_len(arr) + counts * (str_len(new) - str_len(old))
-    out_dtype = f"{arr.dtype.char}{buffersizes.max()}"
-    out = np.empty_like(arr, shape=buffersizes.shape, dtype=out_dtype)
+    arr = np.asanyarray(a)
+    old_dtype = getattr(old, 'dtype', None)
+    old = np.asanyarray(old)
+    new_dtype = getattr(new, 'dtype', None)
+    new = np.asanyarray(new)
+
+    try_out_dt = np.result_type(arr, old, new)
+    if try_out_dt.char == "T":
+        arr = a.astype(try_out_dt, copy=False)
+        old = old.astype(try_out_dt, copy=False)
+        new = new.astype(try_out_dt, copy=False)
+        counts = count
+        out = None
+    else:
+        a_dt = arr.dtype
+        old = old.astype(old_dtype if old_dtype else a_dt, copy=False)
+        new = new.astype(new_dtype if new_dtype else a_dt, copy=False)
+        max_int64 = np.iinfo(np.int64).max
+        counts = _count_ufunc(arr, old, 0, max_int64)
+        counts = np.where(count < 0, counts, np.minimum(counts, count))
+        buffersizes = str_len(arr) + counts * (str_len(new) - str_len(old))
+        out_dtype = f"{arr.dtype.char}{buffersizes.max()}"
+        out = np.empty_like(arr, shape=buffersizes.shape, dtype=out_dtype)
     return _replace(arr, old, new, counts, out=out)
 
 
@@ -1429,11 +1471,15 @@ def partition(a, sep):
 
     """
     a = np.asanyarray(a)
-    # TODO switch to copy=False when issues around views are fixed
-    sep = np.array(sep, dtype=a.dtype, copy=True, subok=True)
-    if a.dtype.char == "T":
+    sep = np.asanyarray(sep)
+
+    try_out_dt = np.result_type(a, sep)
+    if try_out_dt.char == "T":
+        a = a.astype(try_out_dt, copy=False)
+        sep = sep.astype(try_out_dt, copy=False)
         return _partition(a, sep)
 
+    sep = sep.astype(a.dtype, copy=False)
     pos = _find_ufunc(a, sep, 0, MAX)
     a_len = str_len(a)
     sep_len = str_len(sep)
@@ -1495,11 +1541,15 @@ def rpartition(a, sep):
 
     """
     a = np.asanyarray(a)
-    # TODO switch to copy=False when issues around views are fixed
-    sep = np.array(sep, dtype=a.dtype, copy=True, subok=True)
-    if a.dtype.char == "T":
+    sep = np.asanyarray(sep)
+
+    try_out_dt = np.result_type(a, sep)
+    if try_out_dt.char == "T":
+        a = a.astype(try_out_dt, copy=False)
+        sep = sep.astype(try_out_dt, copy=False)
         return _rpartition(a, sep)
 
+    sep = sep.astype(a.dtype, copy=False)
     pos = _rfind_ufunc(a, sep, 0, MAX)
     a_len = str_len(a)
     sep_len = str_len(sep)

numpy/_core/tests/test_stringdtype.py

@@ -1012,13 +1012,16 @@ def test_findlike_promoters():
 
 
 def test_strip_promoter():
-    arg = "Hello!!!!"
+    arg = ["Hello!!!!", "Hello??!!"]
     strip_char = "!"
-    answer = "Hello"
+    answer = ["Hello", "Hello??"]
     for dtypes in [("T", "U"), ("U", "T")]:
-        assert answer == np.strings.strip(
-            np.array(arg, dtype=dtypes[0]), np.array(strip_char, dtype=dtypes[1])
+        result = np.strings.strip(
+            np.array(arg, dtype=dtypes[0]),
+            np.array(strip_char, dtype=dtypes[1])
         )
+        assert_array_equal(result, answer)
+        assert result.dtype.char == "T"
 
 
 def test_replace_promoter():
@@ -1035,15 +1038,18 @@ def test_replace_promoter():
             np.array(new, dtype=dtypes[2]),
         )
         assert_array_equal(answer_arr, answer)
+        assert answer_arr.dtype.char == "T"
 
 
 def test_center_promoter():
-    arg = "Hello, planet!"
+    arg = ["Hello", "planet!"]
     fillchar = "/"
     for dtypes in [("T", "U"), ("U", "T")]:
-        assert "/Hello, planet!/" == np.strings.center(
-            np.array(arg, dtype=dtypes[0]), 16, np.array(fillchar, dtype=dtypes[1])
+        answer = np.strings.center(
+            np.array(arg, dtype=dtypes[0]), 9, np.array(fillchar, dtype=dtypes[1])
         )
+        assert_array_equal(answer, ["//Hello//", "/planet!/"])
+        assert answer.dtype.char == "T"
 
 
 DATETIME_INPUT = [