changes for np2

static-frame · Jul 11, 2024 · 62e905a · 62e905a
1 parent 29bbbdb
commit 62e905a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 12 deletions.
diff --git a/arraymap.c b/arraymap.c
@@ -86,7 +86,8 @@ typedef enum KeysArrayType{
 NPY_DATETIMEUNIT
 dt_unit_from_array(PyArrayObject* a) {
     // This is based on get_datetime_metadata_from_dtype in the NumPy source, but that function is private. This does not check that the dytpe is of the appropriate type.
-    PyArray_DatetimeMetaData* dma = &(((PyArray_DatetimeDTypeMetaData *)PyArray_DESCR(a)->c_metadata)->meta);
+    PyArray_Descr* dt = PyArray_DESCR(a); // borrowed ref
+    PyArray_DatetimeMetaData* dma = &(((PyArray_DatetimeDTypeMetaData *)PyDataType_C_METADATA(dt))->meta);
     return dma->base;
 }
 
@@ -941,7 +942,7 @@ lookup_hash_unicode(
     Py_ssize_t table_pos = hash & mask;
 
     PyArrayObject *a = (PyArrayObject *)self->keys;
-    Py_ssize_t dt_size = PyArray_DESCR(a)->elsize / UCS4_SIZE;
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
     Py_ssize_t cmp_bytes = Py_MIN(key_size, dt_size) * UCS4_SIZE;
 
     Py_hash_t h = 0;
@@ -983,7 +984,7 @@ lookup_hash_string(
     Py_ssize_t table_pos = hash & mask;
 
     PyArrayObject *a = (PyArrayObject *)self->keys;
-    Py_ssize_t dt_size = PyArray_DESCR(a)->elsize;
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
     Py_ssize_t cmp_bytes = Py_MIN(key_size, dt_size);
 
     Py_hash_t h = 0;
@@ -1284,7 +1285,7 @@ lookup_unicode(FAMObject *self, PyObject* key) {
         return -1;
     }
     PyArrayObject *a = (PyArrayObject *)self->keys;
-    Py_ssize_t dt_size = PyArray_DESCR(a)->elsize / UCS4_SIZE;
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
     // if the key_size is greater than the dtype size of the array, we know there cannot be a match
     Py_ssize_t k_size = PyUnicode_GetLength(key);
     if (k_size > dt_size) {
@@ -1305,7 +1306,7 @@ lookup_string(FAMObject *self, PyObject* key) {
         return -1;
     }
     PyArrayObject *a = (PyArrayObject *)self->keys;
-    Py_ssize_t dt_size = PyArray_DESCR(a)->elsize;
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
     Py_ssize_t k_size = PyBytes_GET_SIZE(key);
     if (k_size > dt_size) {
         return -1;
@@ -1650,7 +1651,7 @@ copy_to_new(PyTypeObject *cls, FAMObject *self, FAMObject *new)
     new->key_buffer = NULL;
     if (new->keys_array_type == KAT_UNICODE) {
         PyArrayObject *a = (PyArrayObject *)new->keys;
-        Py_ssize_t dt_size = PyArray_DESCR(a)->elsize / UCS4_SIZE;
+        Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
         new->key_buffer = (Py_UCS4*)PyMem_Malloc((dt_size+1) * UCS4_SIZE);
     }
 
@@ -1831,7 +1832,7 @@ get(FAMObject *self, PyObject *key, PyObject *missing) {
 # define GET_ALL_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func, to_obj_func) \
 {                                                                             \
     char_type* v;                                                             \
-    Py_ssize_t dt_size = PyArray_DESCR(key_array)->elsize / sizeof(char_type);\
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type);\
     Py_ssize_t k_size;                                                        \
     for (; i < key_size; i++) {                                               \
         v = (char_type*)PyArray_GETPTR1(key_array, i);                        \
@@ -2019,7 +2020,7 @@ fam_get_all(FAMObject *self, PyObject *key) {
 # define GET_ANY_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func)    \
 {                                                                             \
     char_type* v;                                                             \
-    Py_ssize_t dt_size = PyArray_DESCR(key_array)->elsize / sizeof(char_type);\
+    Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type);\
     Py_ssize_t k_size;                                                        \
     for (; i < key_size; i++) {                                               \
         v = (char_type*)PyArray_GETPTR1(key_array, i);                        \
@@ -2539,13 +2540,13 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
                 break;
             case KAT_UNICODE: {
                 // Over allocate buffer by 1 so there is room for null at end. This buffer is only used in lookup();
-                Py_ssize_t dt_size = PyArray_DESCR(a)->elsize / UCS4_SIZE;
+                Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
                 fam->key_buffer = (Py_UCS4*)PyMem_Malloc((dt_size+1) * UCS4_SIZE);
                 INSERT_FLEXIBLE(Py_UCS4, insert_unicode, ucs4_get_end_p);
                 break;
             }
             case KAT_STRING: {
-                Py_ssize_t dt_size = PyArray_DESCR(a)->elsize;
+                Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
                 INSERT_FLEXIBLE(char, insert_string, char_get_end_p);
                 break;
             }

diff --git a/setup.py b/setup.py
@@ -22,8 +22,8 @@ def get_ext_dir(*components: tp.Iterable[str]) -> tp.Sequence[str]:
 extension = setuptools.Extension(
     "arraymap",
     ["arraymap.c"],
-    include_dirs=get_ext_dir("numpy", "core", "include"),
-    library_dirs=get_ext_dir("numpy", "core", "lib"),
+    include_dirs=get_ext_dir("numpy", "_core", "include"),
+    library_dirs=get_ext_dir("numpy", "_core", "lib"),
     define_macros=[("AM_VERSION", AM_VERSION)],
     libraries=["npymath"],
 )