Add clear loop implementation

Author: ngoldbaum

stringdtype/meson.build

@@ -43,7 +43,6 @@ py.install_sources(
 py.extension_module(
   '_main',
   srcs,
-  c_args: ['-g', '-O0', '-pg'],
   install: true,
   subdir: 'stringdtype',
   include_directories: includes

stringdtype/pyproject.toml

@@ -28,3 +28,9 @@ dependencies = [
 [tool.ruff]
 line-length = 79
 per-file-ignores = {"__init__.py" = ["F401"]}
+
+[tool.meson-python.args]
+dist = []
+setup = ["-Ddebug=true", "-Doptimization=0"]
+compile = []
+install = []

stringdtype/stringdtype/src/casts.c

@@ -220,9 +220,10 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
         ss *out_ss = ssnewempty(out_num_bytes);
         if (out_ss == NULL) {
             gil_error(PyExc_MemoryError, "ssnewempty failed");
+            return -1;
         }
         char *out_buf = out_ss->buf;
-        for (int i = 0; i < num_codepoints; i++) {
+        for (size_t i = 0; i < num_codepoints; i++) {
             // get code point
             Py_UCS4 code = in[i];
 

stringdtype/stringdtype/src/dtype.c

@@ -20,6 +20,8 @@ new_stringdtype_instance(void)
     new->base.alignment = _Alignof(ss *);
     new->base.flags |= NPY_NEEDS_INIT;
     new->base.flags |= NPY_LIST_PICKLE;
+    new->base.flags |= NPY_ITEM_REFCOUNT;
+    new->base.flags |= NPY_ITEM_IS_POINTER;
 
     return new;
 }
@@ -162,6 +164,36 @@ stringdtype_ensure_canonical(StringDTypeObject *self)
     return self;
 }
 
+static int
+stringdtype_clear_loop(void *NPY_UNUSED(traverse_context),
+                       PyArray_Descr *NPY_UNUSED(descr), char *data,
+                       npy_intp size, npy_intp stride,
+                       NpyAuxData *NPY_UNUSED(auxdata))
+{
+    while (size--) {
+        if (data != NULL) {
+            free(*(ss **)data);
+        }
+        data += stride;
+    }
+
+    return 0;
+}
+
+static int
+stringdtype_get_clear_loop(void *NPY_UNUSED(traverse_context),
+                           PyArray_Descr *NPY_UNUSED(descr),
+                           int NPY_UNUSED(aligned),
+                           npy_intp NPY_UNUSED(fixed_stride),
+                           traverse_loop_function **out_loop,
+                           NpyAuxData **NPY_UNUSED(out_auxdata),
+                           NPY_ARRAYMETHOD_FLAGS *flags)
+{
+    *flags = NPY_METH_REQUIRES_PYAPI | NPY_METH_NO_FLOATINGPOINT_ERRORS;
+    *out_loop = &stringdtype_clear_loop;
+    return 0;
+}
+
 static PyType_Slot StringDType_Slots[] = {
         {NPY_DT_common_instance, &common_instance},
         {NPY_DT_common_dtype, &common_dtype},
@@ -170,6 +202,7 @@ static PyType_Slot StringDType_Slots[] = {
         {NPY_DT_setitem, &stringdtype_setitem},
         {NPY_DT_getitem, &stringdtype_getitem},
         {NPY_DT_ensure_canonical, &stringdtype_ensure_canonical},
+        {NPY_DT_get_clear_loop, &stringdtype_get_clear_loop},
         {0, NULL}};
 
 static PyObject *

stringdtype/stringdtype/src/main.c

@@ -29,9 +29,9 @@ _memory_usage(PyObject *NPY_UNUSED(self), PyObject *obj)
         return NULL;
     }
 
-    NpyIter *iter =
-            NpyIter_New(arr, NPY_ITER_READONLY | NPY_ITER_EXTERNAL_LOOP,
-                        NPY_KEEPORDER, NPY_NO_CASTING, NULL);
+    NpyIter *iter = NpyIter_New(
+            arr, NPY_ITER_READONLY | NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK,
+            NPY_KEEPORDER, NPY_NO_CASTING, NULL);
 
     if (iter == NULL) {
         return NULL;
@@ -90,7 +90,7 @@ PyInit__main(void)
     if (_import_array() < 0) {
         return NULL;
     }
-    if (import_experimental_dtype_api(7) < 0) {
+    if (import_experimental_dtype_api(9) < 0) {
         return NULL;
     }