Redirect NPY_DT_setitem/NPY_DT_getitem to python

Author: trundev

sample_dtypes/scalar.py

@@ -37,15 +37,18 @@ def __repr__(self) -> str:
     return f'{type(self).__name__}(id={hex(id(self))}, shape={self._ndarr.shape}, dtype={self._ndarr.dtype})'
 
   @property
-  def data_buffer(self) -> bytes:
+  def ndarr(self) -> np.ndarray:
+    return self._ndarr
+
+  def setitem(self, target: np.ndarray) -> None:
     """Get raw-data"""
-    return self._ndarr.tobytes(order=BYTES_BUF_ORDER)
+    print(
+        f'TODO: setitem, taget: {type(target)}, {target.shape}, {target.dtype}'
+    )
+    target[...] = self._ndarr
 
-  def from_data_buffer(self, buf: bytes) -> 'SampleScalar':
-    """Create new object using raw-data"""
-    assert (
-        len(buf) == self._ndarr.nbytes
-    ), f'Buffer size must be {self._ndarr.nbytes}'
+  def getitem(self, source: np.ndarray) -> 'SampleScalar':
+    """Store raw-data into a new object"""
     new = self.copy()
-    new._ndarr.flat[:] = np.frombuffer(buf, dtype=self._ndarr.dtype)
+    new._ndarr.flat[:] = source
     return new

src/dtype.c

@@ -103,6 +103,37 @@ sampledtype_ensure_canonical(SampleDTypeObject *self) {
   return self;
 }
 
+static PyObject *create_element_view(SampleDTypeObject *descr, char *dataptr,
+                                     npy_bool readonly) {
+  PyObject *ndarr_obj = PyObject_GetAttrString(descr->sample_scalar, "ndarr");
+  if (ndarr_obj == NULL) {
+    return NULL;
+  }
+  if (!PyArray_Check(ndarr_obj)) {
+    Py_DECREF(ndarr_obj);
+    PyErr_SetString(PyExc_TypeError,
+                    "SampleScalar.ndarr() must return NumPy array");
+    return NULL;
+  }
+
+  PyArrayObject *ndarr_arrobj = (PyArrayObject *)ndarr_obj;
+  PyArray_Descr *ndarr_descr = PyArray_DESCR(ndarr_arrobj);
+
+  PyObject *view = PyArray_NewFromDescr(
+      &PyArray_Type, ndarr_descr, PyArray_NDIM(ndarr_arrobj),
+      PyArray_DIMS(ndarr_arrobj), PyArray_STRIDES(ndarr_arrobj),
+      dataptr, // data pointer for the view
+      readonly ? 0 : NPY_ARRAY_WRITEABLE, NULL);
+  if (view != NULL) {
+    // Now 'view' owns the descriptor, it will be dereferenced when view is
+    // deallocated
+    Py_INCREF(ndarr_descr);
+  }
+
+  Py_DECREF(ndarr_obj);
+  return view;
+}
+
 static int sampledtype_setitem(SampleDTypeObject *descr, PyObject *obj,
                                char *dataptr) {
   printf("%s, target elsise %lld, type_num %d\n", __func__, descr->base.elsize,
@@ -127,29 +158,17 @@ static int sampledtype_setitem(SampleDTypeObject *descr, PyObject *obj,
     return -1;
   }
 
-  // Copy data-bytes from 'obj' into 'dataptr'
-  PyObject *ndarr_bytes = PyObject_GetAttrString(obj, "data_buffer");
-  if (ndarr_bytes == NULL) {
+  PyObject *view = create_element_view(descr, dataptr, NPY_FALSE);
+  if (view == NULL) {
     return -1;
   }
 
-  const char *src_data = NULL;
-  Py_ssize_t data_size = 0;
-  if (PyBytes_AsStringAndSize(ndarr_bytes, &src_data, &data_size) != 0) {
-    Py_DECREF(ndarr_bytes);
-    return -1;
-  }
-
-  if (data_size != descr->base.elsize) {
-    Py_DECREF(ndarr_bytes);
-    PyErr_Format(PyExc_ValueError,
-                 "Unexpected item data buffer size: %zd / %zd", data_size,
-                 descr->base.elsize);
+  res = PyObject_CallMethod(obj, "setitem", "O", view);
+  Py_DECREF(view);
+  if (res == NULL) {
     return -1;
   }
-
-  memcpy(dataptr, src_data, data_size);
-  Py_DECREF(ndarr_bytes);
+  Py_DECREF(res);
 
   return 0;
 }
@@ -158,16 +177,17 @@ static PyObject *sampledtype_getitem(SampleDTypeObject *descr, char *dataptr) {
   printf("%s, source elsize %lld, type_num %d\n", __func__, descr->base.elsize,
          descr->base.type_num);
 
-  // Pass data-bytes from 'dataptr' to 'obj'
-  PyObject *ndarr_bytes =
-      PyBytes_FromStringAndSize(dataptr, descr->base.elsize);
-  if (ndarr_bytes == NULL) {
+  PyObject *view = create_element_view(descr, dataptr, NPY_TRUE);
+  if (view == NULL) {
     return NULL;
   }
 
-  PyObject *res = PyObject_CallMethod(descr->sample_scalar, "from_data_buffer",
-                                      "O", ndarr_bytes);
-  Py_DECREF(ndarr_bytes);
+  PyObject *res =
+      PyObject_CallMethod(descr->sample_scalar, "getitem", "O", view);
+  Py_DECREF(view);
+  if (res == NULL) {
+    return NULL;
+  }
 
   return res;
 }

tests/test_create.py

@@ -1,4 +1,4 @@
-"""Test basic module import"""
+"""Test DType creation/instantiation"""
 
 import numpy as np
 import pytest
@@ -31,12 +31,12 @@ def test_setitem():
   scalar = sample_dtypes.SampleScalar(dtype=int)
 
   print('  * set compatible value')
-  arr = np.empty(3, sample_dtypes.SampleDType(scalar))
+  arr: np.ndarray = np.empty(3, sample_dtypes.SampleDType(scalar))
   scalar._ndarr.flat[...] = np.arange(scalar._ndarr.size) + 10
   arr[1] = scalar
   np.testing.assert_equal(arr[1]._ndarr, scalar._ndarr)
 
   print('  * set incompatible value')
-  arr: np.ndarray = np.empty(3, sample_dtypes.SampleDType)
+  arr = np.empty(3, sample_dtypes.SampleDType)
   with pytest.raises(ValueError):
     arr[0] = scalar

tests/test_refcount.py

@@ -0,0 +1,25 @@
+"""Test internal object reference counting"""
+
+import sys
+import numpy as np
+
+
+def test_dtype_refcnt():
+  """Test SampleScalar internal dtype reference counts
+
+  NOTE: This fails when debugging as of extra ownership of local object
+  """
+  import sample_dtypes
+
+  # Use non-singleton (non-immortal) dtypes to have valid ref-counts
+  dtype = np.dtype([('a', np.int16)])
+  assert (
+      sys.getrefcount(dtype) == 2
+  ), 'Unexpected initial ref-count for non-immortal object'
+
+  scalar = sample_dtypes.SampleScalar(dtype=dtype)
+  arr: np.ndarray = np.empty(3, sample_dtypes.SampleDType(scalar))
+  arr[1] = scalar
+
+  del arr, scalar
+  assert sys.getrefcount(dtype) == 2, 'Unexpected dtype ref-count'