ngoldbaum · Mar 12, 2024 · Feb 13, 2024 · Mar 1, 2024 · Mar 1, 2024 · Mar 1, 2024
diff --git a/numpy/_core/code_generators/generate_umath.py b/numpy/_core/code_generators/generate_umath.py
@@ -1280,7 +1280,26 @@ def english_upper(s):
           docstrings.get('numpy._core.umath._expandtabs'),
           None,
           ),
-
+'_center':
+    Ufunc(3, 1, None,
+          docstrings.get('numpy._core.umath._center'),
+          None,
+          ),
+'_ljust':
+    Ufunc(3, 1, None,
+          docstrings.get('numpy._core.umath._ljust'),
+          None,
+          ),
+'_rjust':
+    Ufunc(3, 1, None,
+          docstrings.get('numpy._core.umath._rjust'),
+          None,
+          ),
+'_zfill':
+    Ufunc(2, 1, None,
+          docstrings.get('numpy._core.umath._zfill'),
+          None,
+          ),
 }
 
 def indent(st, spaces):

diff --git a/numpy/_core/code_generators/ufunc_docstrings.py b/numpy/_core/code_generators/ufunc_docstrings.py
@@ -4885,3 +4885,143 @@ def add_newdoc(place, name, doc):
 
 add_newdoc('numpy._core.umath', '_expandtabs_length', '')
 add_newdoc('numpy._core.umath', '_expandtabs', '')
+
+add_newdoc('numpy._core.umath', '_center',
+    """
+    Return a copy of `x1` with its elements centered in a string of
+    length `x2`.
+
+    Parameters
+    ----------
+    x1 : array_like, with `np.bytes_` or `np.str_` dtype
+
+    x2 : array_like, with any integer dtype
+        The length of the resulting strings, unless ``width < str_len(a)``.
+    x3 : array_like, with `np.bytes_` or `np.str_` dtype
+        The padding character to use.
+        $PARAMS
+
+    Returns
+    -------
+    out : ndarray
+        Output array of str or unicode, depending on input
+        types
+        $OUT_SCALAR_2
+
+    See Also
+    --------
+    str.center
+
+    Examples
+    --------
+    >>> c = np.array(['a1b2','1b2a','b2a1','2a1b']); c
+    array(['a1b2', '1b2a', 'b2a1', '2a1b'], dtype='<U4')
+    >>> np.strings.center(c, width=9)
+    array(['   a1b2  ', '   1b2a  ', '   b2a1  ', '   2a1b  '], dtype='<U9')
+    >>> np.strings.center(c, width=9, fillchar='*')
+    array(['***a1b2**', '***1b2a**', '***b2a1**', '***2a1b**'], dtype='<U9')
+    >>> np.strings.center(c, width=1)
+    array(['a1b2', '1b2a', 'b2a1', '2a1b'], dtype='<U4')
+
+    """)
+
+add_newdoc('numpy._core.umath', '_ljust',
+    """
+    Return an array with the elements of `x1` left-justified in a
+    string of length `x2`.
+
+    Parameters
+    ----------
+    x1 : array_like, with `np.bytes_` or `np.str_` dtype
+
+    x2 : array_like, with any integer dtype
+        The length of the resulting strings, unless ``width < str_len(a)``.
+    x3 : array_like, with `np.bytes_` or `np.str_` dtype
+        The padding character to use.
+        $PARAMS
+
+    Returns
+    -------
+    out : ndarray
+        Output array of str or unicode, depending on input type
+        $OUT_SCALAR_2
+
+    See Also
+    --------
+    str.ljust
+
+    Examples
+    --------
+    >>> c = np.array(['aAaAaA', '  aA  ', 'abBABba'])
+    >>> np.strings.ljust(c, width=3)
+    array(['aAaAaA', '  aA  ', 'abBABba'], dtype='<U7')
+    >>> np.strings.ljust(c, width=9)
+    array(['aAaAaA   ', '  aA     ', 'abBABba  '], dtype='<U9')
+
+    """)
+
+add_newdoc('numpy._core.umath', '_rjust',
+    """
+    Return an array with the elements of `x1` right-justified in a
+    string of length `x2`.
+
+    Parameters
+    ----------
+    x1 : array_like, with `np.bytes_` or `np.str_` dtype
+
+    x2 : array_like, with any integer dtype
+        The length of the resulting strings, unless ``width < str_len(a)``.
+    x3 : array_like, with `np.bytes_` or `np.str_` dtype
+        The padding character to use.
+        $PARAMS
+
+    Returns
+    -------
+    out : ndarray
+        Output array of str or unicode, depending on input type
+        $OUT_SCALAR_2
+
+    See Also
+    --------
+    str.rjust
+
+    Examples
+    --------
+    >>> a = np.array(['aAaAaA', '  aA  ', 'abBABba'])
+    >>> np.strings.rjust(a, width=3)
+    array(['aAaAaA', '  aA  ', 'abBABba'], dtype='<U7')
+    >>> np.strings.rjust(a, width=9)
+    array(['   aAaAaA', '     aA  ', '  abBABba'], dtype='<U9')
+
+    """)
+
+add_newdoc('numpy._core.umath', '_zfill',
+    """
+    Return the numeric string left-filled with zeros. A leading
+    sign prefix (``+``/``-``) is handled by inserting the padding
+    after the sign character rather than before.
+
+    Parameters
+    ----------
+    x1 : array_like, with `np.bytes_` or `np.str_` dtype
+
+    x2 : array_like, with any integer dtype
+        Width of string to left-fill elements in `a`.
+        $PARAMS
+
+    Returns
+    -------
+    out : ndarray
+        Output array of str or unicode, depending on input type
+        $OUT_SCALAR_2
+
+    See Also
+    --------
+    str.zfill
+
+    Examples
+    --------
+    >>> np.strings.zfill(['1', '-1', '+1'], 3)
+    array(['001', '-01', '+01'], dtype='<U3')
+
+    """)
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.c b/numpy/_core/src/multiarray/stringdtype/dtype.c
@@ -922,8 +922,7 @@ load_new_string(npy_packed_static_string *out, npy_static_string *out_ss,
                       "Failed to allocate string in %s", err_context);
         return -1;
     }
-    int is_null = NpyString_load(allocator, out_pss, out_ss);
-    if (is_null == -1) {
+    if (NpyString_load(allocator, out_pss, out_ss) == -1) {
         npy_gil_error(PyExc_MemoryError,
                       "Failed to load string in %s", err_context);
         return -1;

diff --git a/numpy/_core/src/umath/string_buffer.h b/numpy/_core/src/umath/string_buffer.h
@@ -1505,4 +1505,92 @@ string_expandtabs(Buffer<enc> buf, npy_int64 tabsize, Buffer<enc> out)
 }
 
 
+enum class JUSTPOSITION {
+    CENTER, LEFT, RIGHT
+};
+
+template <ENCODING enc>
+static inline npy_intp
+string_pad(Buffer<enc> buf, npy_int64 width, npy_ucs4 fill, JUSTPOSITION pos, Buffer<enc> out)
+{
+    size_t finalwidth = width > 0 ? width : 0;
+    if (finalwidth > PY_SSIZE_T_MAX) {
+        npy_gil_error(PyExc_OverflowError, "padded string is too long");
+        return -1;
+    }
+
+    size_t len_codepoints = buf.num_codepoints();
+    size_t len_bytes = buf.after - buf.buf;
+
+    size_t len;
+    if (enc == ENCODING::UTF8) {
+        len = len_bytes;
+    }
+    else {
+        len = len_codepoints;
+    }
+
+    if (len_codepoints >= finalwidth) {
+        buf.buffer_memcpy(out, len);
+        return (npy_intp) len;
+    }
+
+    size_t left, right;
+    if (pos == JUSTPOSITION::CENTER) {
+        size_t pad = finalwidth - len_codepoints;
+        left = pad / 2 + (pad & finalwidth & 1);
+        right = pad - left;
+    }
+    else if (pos == JUSTPOSITION::LEFT) {
+        left = 0;
+        right = finalwidth - len_codepoints;
+    }
+    else {
+        left = finalwidth - len_codepoints;
+        right = 0;
+    }
+
+    assert(left >= 0 || right >= 0);
+    assert(left <= PY_SSIZE_T_MAX - len && right <= PY_SSIZE_T_MAX - (left + len));
+
+    if (left > 0) {
+        out.advance_chars_or_bytes(out.buffer_memset(fill, left));
+    }
+
+    buf.buffer_memcpy(out, len);
+    out += len_codepoints;
+
+    if (right > 0) {
+        out.advance_chars_or_bytes(out.buffer_memset(fill, right));
+    }
+
+    return finalwidth;
+}
+
+
+template <ENCODING enc>
+static inline npy_intp
+string_zfill(Buffer<enc> buf, npy_int64 width, Buffer<enc> out)
+{
+    size_t finalwidth = width > 0 ? width : 0;
+
+    npy_ucs4 fill = '0';
+    npy_intp new_len = string_pad(buf, width, fill, JUSTPOSITION::RIGHT, out);
+    if (new_len == -1) {
+        return -1;
+    }
+
+    size_t offset = finalwidth - buf.num_codepoints();
+    Buffer<enc> tmp = out + offset;
+
+    npy_ucs4 c = *tmp;
+    if (c == '+' || c == '-') {
+        tmp.buffer_memset(fill, 1);
+        out.buffer_memset(c, 1);
+    }
+
+    return new_len;
+}
+
+
 #endif /* _NPY_CORE_SRC_UMATH_STRING_BUFFER_H_ */