Merge pull request numpy#25891 from lysnikolaou/string-ufuncs-expandtabs

ENH: Add expandtabs ufunc for string & unicode dtypes

Author: ngoldbaum

numpy/_core/code_generators/generate_umath.py

@@ -1270,6 +1270,17 @@ def english_upper(s):
           docstrings.get('numpy._core.umath._rstrip_whitespace'),
           None,
           ),
+'_expandtabs_length':
+    Ufunc(2, 1, None,
+          docstrings.get('numpy._core.umath._expandtabs_length'),
+          None,
+          ),
+'_expandtabs':
+    Ufunc(2, 1, None,
+          docstrings.get('numpy._core.umath._expandtabs'),
+          None,
+          ),
+
 }
 
 def indent(st, spaces):

numpy/_core/code_generators/ufunc_docstrings.py

@@ -4882,3 +4882,6 @@ def add_newdoc(place, name, doc):
 add_newdoc('numpy._core.umath', '_strip_whitespace', '')
 add_newdoc('numpy._core.umath', '_lstrip_whitespace', '')
 add_newdoc('numpy._core.umath', '_rstrip_whitespace', '')
+
+add_newdoc('numpy._core.umath', '_expandtabs_length', '')
+add_newdoc('numpy._core.umath', '_expandtabs', '')

numpy/_core/src/umath/string_buffer.h

@@ -404,28 +404,36 @@ struct Buffer {
         }
     }
 
-    inline void
+    inline npy_intp
     buffer_memset(npy_ucs4 fill_char, size_t n_chars)
     {
         if (n_chars == 0) {
-            return;
+            return 0;
         }
         switch (enc) {
             case ENCODING::ASCII:
-                memset(buf, fill_char, n_chars);
-                break;
+                memset(this->buf, fill_char, n_chars);
+                return n_chars;
             case ENCODING::UTF32:
             {
-                char *tmp = buf;
+                char *tmp = this->buf;
                 for (size_t i = 0; i < n_chars; i++) {
                     *(npy_ucs4 *)tmp = fill_char;
                     tmp += sizeof(npy_ucs4);
                 }
-                break;
+                return n_chars;
             }
             case ENCODING::UTF8:
-                assert(false);  // buffer_memset not used by stringdtype
-                break;
+            {
+                char utf8_c[4] = {0};
+                char *tmp = this->buf;
+                size_t num_bytes = ucs4_code_to_utf8_char(fill_char, utf8_c);
+                for (size_t i = 0; i < n_chars; i++) {
+                    memcpy(tmp, utf8_c, num_bytes);
+                    tmp += num_bytes;
+                }
+                return num_bytes * n_chars;
+            }
         }
     }
 
@@ -1428,5 +1436,73 @@ string_replace(Buffer<enc> buf1, Buffer<enc> buf2, Buffer<enc> buf3, npy_int64 c
 }
 
 
+template <ENCODING enc>
+static inline npy_intp
+string_expandtabs_length(Buffer<enc> buf, npy_int64 tabsize)
+{
+    size_t len = buf.num_codepoints();
+
+    npy_intp new_len = 0, line_pos = 0;
+
+    Buffer<enc> tmp = buf;
+    for (size_t i = 0; i < len; i++) {
+        npy_ucs4 ch = *tmp;
+        if (ch == '\t') {
+            if (tabsize > 0) {
+                npy_intp incr = tabsize - (line_pos % tabsize);
+                line_pos += incr;
+                new_len += incr;
+            }
+        }
+        else {
+            line_pos += 1;
+            size_t n_bytes = tmp.num_bytes_next_character();
+            new_len += n_bytes;
+            if (ch == '\n' || ch == '\r') {
+                line_pos = 0;
+            }
+        }
+        if (new_len == PY_SSIZE_T_MAX || new_len < 0) {
+            npy_gil_error(PyExc_OverflowError, "new string is too long");
+            return -1;
+        }
+        tmp++;
+    }
+    return new_len;
+}
+
+
+template <ENCODING enc>
+static inline npy_intp
+string_expandtabs(Buffer<enc> buf, npy_int64 tabsize, Buffer<enc> out)
+{
+    size_t len = buf.num_codepoints();
+
+    npy_intp new_len = 0, line_pos = 0;
+
+    Buffer<enc> tmp = buf;
+    for (size_t i = 0; i < len; i++) {
+        npy_ucs4 ch = *tmp;
+        if (ch == '\t') {
+            if (tabsize > 0) {
+                npy_intp incr = tabsize - (line_pos % tabsize);
+                line_pos += incr;
+                new_len += out.buffer_memset((npy_ucs4) ' ', incr);
+                out += incr;
+            }
+        }
+        else {
+            line_pos++;
+            new_len += out.buffer_memset(ch, 1);
+            out++;
+            if (ch == '\n' || ch == '\r') {
+                line_pos = 0;
+            }
+        }
+        tmp++;
+    }
+    return new_len;
+}
+
 
 #endif /* _NPY_CORE_SRC_UMATH_STRING_BUFFER_H_ */

numpy/_core/src/umath/string_ufuncs.cpp

@@ -450,6 +450,63 @@ string_lrstrip_chars_loop(PyArrayMethod_Context *context,
 }
 
 
+template <ENCODING enc>
+static int
+string_expandtabs_length_loop(PyArrayMethod_Context *context,
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    int elsize = context->descriptors[0]->elsize;
+
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+
+    npy_intp N = dimensions[0];
+
+    while (N--) {
+        Buffer<enc> buf(in1, elsize);
+        *(npy_intp *)out = string_expandtabs_length(buf, *(npy_int64 *)in2);
+
+        in1 += strides[0];
+        in2 += strides[1];
+        out += strides[2];
+    }
+
+    return 0;
+}
+
+
+template <ENCODING enc>
+static int
+string_expandtabs_loop(PyArrayMethod_Context *context,
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    int elsize = context->descriptors[0]->elsize;
+    int outsize = context->descriptors[2]->elsize;
+
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+
+    npy_intp N = dimensions[0];
+
+    while (N--) {
+        Buffer<enc> buf(in1, elsize);
+        Buffer<enc> outbuf(out, outsize);
+        npy_intp new_len = string_expandtabs(buf, *(npy_int64 *)in2, outbuf);
+        outbuf.buffer_fill_with_zeros_after_index(new_len);
+
+        in1 += strides[0];
+        in2 += strides[1];
+        out += strides[2];
+    }
+
+    return 0;
+}
+
+
 /* Resolve descriptors & promoter functions */
 
 static NPY_CASTING
@@ -517,9 +574,9 @@ string_multiply_resolve_descriptors(
 static NPY_CASTING
 string_strip_whitespace_resolve_descriptors(
         PyArrayMethodObject *NPY_UNUSED(self),
-        PyArray_DTypeMeta *NPY_UNUSED(dtypes[3]),
-        PyArray_Descr *given_descrs[3],
-        PyArray_Descr *loop_descrs[3],
+        PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
+        PyArray_Descr *given_descrs[2],
+        PyArray_Descr *loop_descrs[2],
         npy_intp *NPY_UNUSED(view_offset))
 {
     loop_descrs[0] = NPY_DT_CALL_ensure_canonical(given_descrs[0]);
@@ -600,7 +657,7 @@ string_replace_promoter(PyObject *NPY_UNUSED(ufunc),
 static NPY_CASTING
 string_replace_resolve_descriptors(
         PyArrayMethodObject *NPY_UNUSED(self),
-        PyArray_DTypeMeta *NPY_UNUSED(dtypes[3]),
+        PyArray_DTypeMeta *NPY_UNUSED(dtypes[5]),
         PyArray_Descr *given_descrs[5],
         PyArray_Descr *loop_descrs[5],
         npy_intp *NPY_UNUSED(view_offset))
@@ -651,6 +708,67 @@ string_startswith_endswith_promoter(PyObject *NPY_UNUSED(ufunc),
 }
 
 
+static int
+string_expandtabs_length_promoter(PyObject *NPY_UNUSED(ufunc),
+        PyArray_DTypeMeta *op_dtypes[], PyArray_DTypeMeta *signature[],
+        PyArray_DTypeMeta *new_op_dtypes[])
+{
+    Py_INCREF(op_dtypes[0]);
+    new_op_dtypes[0] = op_dtypes[0];
+    new_op_dtypes[1] = NPY_DT_NewRef(&PyArray_Int64DType);
+    new_op_dtypes[2] = PyArray_DTypeFromTypeNum(NPY_DEFAULT_INT);
+    return 0;
+}
+
+
+static int
+string_expandtabs_promoter(PyObject *NPY_UNUSED(ufunc),
+        PyArray_DTypeMeta *op_dtypes[], PyArray_DTypeMeta *signature[],
+        PyArray_DTypeMeta *new_op_dtypes[])
+{
+    Py_INCREF(op_dtypes[0]);
+    new_op_dtypes[0] = op_dtypes[0];
+    new_op_dtypes[1] = NPY_DT_NewRef(&PyArray_Int64DType);
+    Py_INCREF(op_dtypes[0]);
+    new_op_dtypes[2] = op_dtypes[0];
+    return 0;
+}
+
+
+static NPY_CASTING
+string_expandtabs_resolve_descriptors(
+        PyArrayMethodObject *NPY_UNUSED(self),
+        PyArray_DTypeMeta *NPY_UNUSED(dtypes[3]),
+        PyArray_Descr *given_descrs[3],
+        PyArray_Descr *loop_descrs[3],
+        npy_intp *NPY_UNUSED(view_offset))
+{
+    if (given_descrs[2] == NULL) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "The 'out' kwarg is necessary. Use numpy.strings.expandtabs without it.");
+        return _NPY_ERROR_OCCURRED_IN_CAST;
+    }
+
+    loop_descrs[0] = NPY_DT_CALL_ensure_canonical(given_descrs[0]);
+    if (loop_descrs[0] == NULL) {
+        return _NPY_ERROR_OCCURRED_IN_CAST;
+    }
+
+    loop_descrs[1] = NPY_DT_CALL_ensure_canonical(given_descrs[1]);
+    if (loop_descrs[1] == NULL) {
+        return _NPY_ERROR_OCCURRED_IN_CAST;
+    }
+
+    loop_descrs[2] = NPY_DT_CALL_ensure_canonical(given_descrs[2]);
+    if (loop_descrs[2] == NULL) {
+        return _NPY_ERROR_OCCURRED_IN_CAST;
+    }
+
+    return NPY_NO_CASTING;
+}
+
+
 /*
  * Machinery to add the string loops to the existing ufuncs.
  */
@@ -1130,6 +1248,42 @@ init_string_ufuncs(PyObject *umath)
         }
     }
 
+    dtypes[0] = NPY_OBJECT;
+    dtypes[1] = NPY_INT64;
+    dtypes[2] = NPY_DEFAULT_INT;
+    if (init_ufunc(
+            umath, "_expandtabs_length", 2, 1, dtypes, ENCODING::ASCII,
+            string_expandtabs_length_loop<ENCODING::ASCII>, NULL, NULL) < 0) {
+        return -1;
+    }
+    if (init_ufunc(
+            umath, "_expandtabs_length", 2, 1, dtypes, ENCODING::UTF32,
+            string_expandtabs_length_loop<ENCODING::UTF32>, NULL, NULL) < 0) {
+        return -1;
+    }
+    if (init_promoter(umath, "_expandtabs_length", 2, 1, string_expandtabs_length_promoter) < 0) {
+        return -1;
+    }
+
+    dtypes[0] = NPY_OBJECT;
+    dtypes[1] = NPY_INT64;
+    dtypes[2] = NPY_OBJECT;
+    if (init_ufunc(
+            umath, "_expandtabs", 2, 1, dtypes, ENCODING::ASCII,
+            string_expandtabs_loop<ENCODING::ASCII>,
+            string_expandtabs_resolve_descriptors, NULL) < 0) {
+        return -1;
+    }
+    if (init_ufunc(
+            umath, "_expandtabs", 2, 1, dtypes, ENCODING::UTF32,
+            string_expandtabs_loop<ENCODING::UTF32>,
+            string_expandtabs_resolve_descriptors, NULL) < 0) {
+        return -1;
+    }
+    if (init_promoter(umath, "_expandtabs", 2, 1, string_expandtabs_promoter) < 0) {
+        return -1;
+    }
+
     return 0;
 }