Merge pull request #85 from ngoldbaum/alloc-refactor

Use PyMem allocators
numpy · Aug 30, 2023 · fb1b5ba · fb1b5ba
2 parents 018ab87 + 58ac555
commit fb1b5ba
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 14 deletions.
diff --git a/asciidtype/tests/test_asciidtype.py b/asciidtype/tests/test_asciidtype.py
@@ -157,7 +157,7 @@ def test_casting_safety():
 def test_unicode_to_ascii_to_unicode():
     arr = np.array(["hello", "this", "is", "an", "array"])
     ascii_arr = arr.astype(ASCIIDType(5))
-    for dtype in ["U5", np.unicode_, np.str_]:
+    for dtype in ["U5", np.str_]:
         round_trip_arr = ascii_arr.astype(dtype)
         np.testing.assert_array_equal(arr, round_trip_arr)
 

diff --git a/stringdtype/stringdtype/src/casts.c b/stringdtype/stringdtype/src/casts.c
@@ -1171,7 +1171,7 @@ get_cast_spec(const char *name, NPY_CASTING casting,
               NPY_ARRAYMETHOD_FLAGS flags, PyArray_DTypeMeta **dtypes,
               PyType_Slot *slots)
 {
-    PyArrayMethod_Spec *ret = malloc(sizeof(PyArrayMethod_Spec));
+    PyArrayMethod_Spec *ret = PyMem_Malloc(sizeof(PyArrayMethod_Spec));
 
     ret->name = name;
     ret->nin = 1;
@@ -1187,7 +1187,7 @@ get_cast_spec(const char *name, NPY_CASTING casting,
 PyArray_DTypeMeta **
 get_dtypes(PyArray_DTypeMeta *dt1, PyArray_DTypeMeta *dt2)
 {
-    PyArray_DTypeMeta **ret = malloc(2 * sizeof(PyArray_DTypeMeta *));
+    PyArray_DTypeMeta **ret = PyMem_Malloc(2 * sizeof(PyArray_DTypeMeta *));
 
     ret[0] = dt1;
     ret[1] = dt2;
@@ -1297,7 +1297,7 @@ get_casts()
             dt2s_dtypes, dt2s_slots);
 
     PyArrayMethod_Spec **casts =
-            malloc((num_casts + 1) * sizeof(PyArrayMethod_Spec *));
+            PyMem_Malloc((num_casts + 1) * sizeof(PyArrayMethod_Spec *));
 
     int cast_i = 0;
 

diff --git a/stringdtype/stringdtype/src/dtype.c b/stringdtype/stringdtype/src/dtype.c
@@ -788,10 +788,12 @@ init_string_dtype(void)
     StringDType.base.singleton = singleton;
 
     for (int i = 0; StringDType_casts[i] != NULL; i++) {
-        free(StringDType_casts[i]->dtypes);
-        free(StringDType_casts[i]);
+        PyMem_Free(StringDType_casts[i]->dtypes);
+        PyMem_Free(StringDType_casts[i]);
     }
 
+    PyMem_Free(StringDType_casts);
+
     return 0;
 }
 

diff --git a/stringdtype/stringdtype/src/static_string.c b/stringdtype/stringdtype/src/static_string.c
@@ -1,3 +1,5 @@
+#include "Python.h"
+
 #include "static_string.h"
 
 // defined this way so it has an in-memory representation that is distinct
@@ -16,7 +18,7 @@ ssnewlen(const char *init, size_t len, ss *to_init)
         return 0;
     }
 
-    char *ret_buf = (char *)malloc(sizeof(char) * len);
+    char *ret_buf = (char *)PyMem_RawMalloc(sizeof(char) * len);
 
     if (ret_buf == NULL) {
         return -1;
@@ -35,7 +37,7 @@ void
 ssfree(ss *str)
 {
     if (str->buf != NULL && str->buf != EMPTY_STRING.buf) {
-        free(str->buf);
+        PyMem_RawFree(str->buf);
         str->buf = NULL;
     }
     str->len = 0;
@@ -68,7 +70,7 @@ ssnewemptylen(size_t num_bytes, ss *out)
         return 0;
     }
 
-    char *buf = (char *)malloc(sizeof(char) * num_bytes);
+    char *buf = (char *)PyMem_RawMalloc(sizeof(char) * num_bytes);
 
     if (buf == NULL) {
         return -1;

diff --git a/stringdtype/stringdtype/src/static_string.h b/stringdtype/stringdtype/src/static_string.h
@@ -22,9 +22,10 @@ ssnewlen(const char *init, size_t len, ss *to_init);
 void
 ssfree(ss *str);
 
-// copies the contents out *in* into *out*. Allocates a new string buffer for
-// *out*. Returns -1 if malloc fails and -2 if *out* is not empty. Returns 0 on
-// success.
+// Copies the contents of *in* into *out*. Allocates a new string buffer for
+// *out* and assumes that *out* is uninitialized. Returns -1 if malloc fails
+// and -2 if *out* is not initialized. ssfree *must* be called before this is
+// called if *in* points to an existing string. Returns 0 on success.
 int
 ssdup(const ss *in, ss *out);
 

diff --git a/stringdtype/tests/test_char.py b/stringdtype/tests/test_char.py
@@ -14,7 +14,7 @@ def string_array():
 
 @pytest.fixture
 def unicode_array():
-    return np.array(TEST_DATA, dtype=np.unicode_)
+    return np.array(TEST_DATA, dtype=np.str_)
 
 
 UNARY_FUNCTIONS = [

diff --git a/stringdtype/tests/test_stringdtype.py b/stringdtype/tests/test_stringdtype.py
@@ -145,7 +145,7 @@ def test_scalars_string_conversion(data, dtype):
     ],
 )
 def test_unicode_casts(dtype, strings):
-    arr = np.array(strings, dtype=np.unicode_).astype(dtype)
+    arr = np.array(strings, dtype=np.str_).astype(dtype)
     expected = np.array(strings, dtype=dtype)
     np.testing.assert_array_equal(arr, expected)