Merge pull request #36 from eriknw/support_udts

Update Cython code to support UDTs

Author: eriknw

suitesparse_graphblas/utils.pxd

@@ -1,16 +1,20 @@
 from libc.stdint cimport uint64_t
-from numpy cimport ndarray
+from numpy cimport ndarray, npy_intp, dtype as dtype_t
 
 
 cdef extern from "numpy/arrayobject.h" nogil:
     # These aren't public (i.e., "extern"), but other projects use them too
-    void *PyDataMem_NEW(size_t)
-    void *PyDataMem_NEW_ZEROED(size_t, size_t)
-    void *PyDataMem_RENEW(void *, size_t)
-    void PyDataMem_FREE(void *)
+    void *PyDataMem_NEW(size_t size)
+    void *PyDataMem_NEW_ZEROED(size_t nmemb, size_t size)
+    void *PyDataMem_RENEW(void *ptr, size_t size)
+    void PyDataMem_FREE(void *ptr)
     # These are available in newer Cython versions
-    void PyArray_ENABLEFLAGS(ndarray, int flags)
-    void PyArray_CLEARFLAGS(ndarray, int flags)
+    void PyArray_ENABLEFLAGS(ndarray array, int flags)
+    void PyArray_CLEARFLAGS(ndarray array, int flags)
+    # Not exposed by Cython (b/c it steals a reference from dtype)
+    ndarray PyArray_NewFromDescr(
+        type subtype, dtype_t dtype, int nd, npy_intp *dims, npy_intp *strides, void *data, int flags, object obj
+    )
 
 ctypedef enum GrB_Mode:
     GrB_NONBLOCKING
@@ -24,11 +28,13 @@ ctypedef uint64_t (*GxB_init)(
     void (*user_free_function)(void *),
 )
 
-cpdef int call_gxb_init(ffi, lib, int mode)
+cpdef int call_gxb_init(object ffi, object lib, int mode)
 
-cpdef ndarray claim_buffer(ffi, cdata, size_t size, dtype)
+cpdef ndarray claim_buffer(object ffi, object cdata, size_t size, dtype_t dtype)
 
-cpdef ndarray claim_buffer_2d(ffi, cdata, size_t cdata_size, size_t nrows, size_t ncols, dtype, bint is_c_order)
+cpdef ndarray claim_buffer_2d(
+    object ffi, object cdata, size_t cdata_size, size_t nrows, size_t ncols, dtype_t dtype, bint is_c_order
+)
 
 cpdef unclaim_buffer(ndarray array)
 

suitesparse_graphblas/utils.pyx

@@ -1,19 +1,19 @@
 import numpy as np
+from cpython.ref cimport Py_INCREF
 from libc.stdint cimport uintptr_t
 from numpy cimport (
     NPY_ARRAY_F_CONTIGUOUS,
     NPY_ARRAY_OWNDATA,
     NPY_ARRAY_WRITEABLE,
-    PyArray_New,
-    PyArray_SimpleNewFromData,
     import_array,
     ndarray,
     npy_intp,
+    dtype as dtype_t,
 )
 
 import_array()
 
-cpdef int call_gxb_init(ffi, lib, int mode):
+cpdef int call_gxb_init(object ffi, object lib, int mode):
     # We need to call `GxB_init`, but we didn't compile Cython against GraphBLAS.  So, we get it from cffi.
     # Step 1: ffi.addressof(lib, "GxB_init")
     #    Return type: cffi.cdata object of a function pointer.  Can't cast to int.
@@ -30,32 +30,37 @@ cpdef int call_gxb_init(ffi, lib, int mode):
     return func(<GrB_Mode>mode, PyDataMem_NEW, PyDataMem_NEW_ZEROED, PyDataMem_RENEW, PyDataMem_FREE)
 
 
-cpdef ndarray claim_buffer(ffi, cdata, size_t size, dtype):
+cpdef ndarray claim_buffer(object ffi, object cdata, size_t size, dtype_t dtype):
     cdef:
         npy_intp dims = size
         uintptr_t ptr = int(ffi.cast("uintptr_t", cdata))
-        ndarray array = PyArray_SimpleNewFromData(1, &dims, dtype.num, <void*>ptr)
+        ndarray array
+    Py_INCREF(dtype)
+    array = PyArray_NewFromDescr(
+        ndarray, dtype, 1, &dims, NULL, <void*>ptr, NPY_ARRAY_WRITEABLE, <object>NULL
+    )
     PyArray_ENABLEFLAGS(array, NPY_ARRAY_OWNDATA)
     return array
 
 
-cpdef ndarray claim_buffer_2d(ffi, cdata, size_t cdata_size, size_t nrows, size_t ncols, dtype, bint is_c_order):
+cpdef ndarray claim_buffer_2d(
+    object ffi, object cdata, size_t cdata_size, size_t nrows, size_t ncols, dtype_t dtype, bint is_c_order
+):
     cdef:
         size_t size = nrows * ncols
         ndarray array
         uintptr_t ptr
         npy_intp dims[2]
+        int flags = NPY_ARRAY_WRITEABLE
     if cdata_size == size:
         ptr = int(ffi.cast("uintptr_t", cdata))
         dims[0] = nrows
         dims[1] = ncols
-        if is_c_order:
-            array = PyArray_SimpleNewFromData(2, dims, dtype.num, <void*>ptr)
-        else:
-            array = PyArray_New(
-                ndarray, 2, dims, dtype.num, NULL, <void*>ptr, -1,
-                NPY_ARRAY_F_CONTIGUOUS | NPY_ARRAY_WRITEABLE, <object>NULL
-            )
+        if not is_c_order:
+            flags |= NPY_ARRAY_F_CONTIGUOUS
+        array = PyArray_NewFromDescr(
+            ndarray, dtype, 2, dims, NULL, <void*>ptr, flags, <object>NULL
+        )
         PyArray_ENABLEFLAGS(array, NPY_ARRAY_OWNDATA)
     elif cdata_size > size:  # pragma: no cover
         array = claim_buffer(ffi, cdata, cdata_size, dtype)