BUG: fix more issues with string ufunc promotion

numpy · charris · Oct 30, 2024 · Oct 24, 2024 · Oct 25, 2024 · Oct 26, 2024
commit a90fe7cc6e52be410321b5fffe4e85e76acedb0a
diff --git a/numpy/_core/strings.py b/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)


@@ -1421,11 +1463,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)
@@ -1487,11 +1533,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)

diff --git a/numpy/_core/tests/test_stringdtype.py b/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 = [