Pure python implementation of NPY_DT_setitem/NPY_DT_getitem

Author: trundev

sample_dtypes/scalar.py

@@ -1,5 +1,6 @@
 """Scalar types needed by the dtype machinery."""
 
+import ctypes
 import numpy as np
 
 
@@ -26,29 +27,35 @@ def alignment(self) -> int:
     return self._ndarr.dtype.alignment
 
   def is_compatible(self, other: 'SampleScalar') -> bool:
-    print(f'TODO: scalar.py: SampleScalar.is_compatible({other})')
-    # This implies `self.elsize == other.elsize and self.alignment == other.alignment`
-    return (
-        self._ndarr.shape == other._ndarr.shape
-        and self._ndarr.dtype == other._ndarr.dtype
-    )
+    print(f'scalar.py: SampleScalar.is_compatible({other})')
+    return self._ndarr.shape == other._ndarr.shape
 
   def __repr__(self) -> str:
     return f'{type(self).__name__}(id={hex(id(self))}, shape={self._ndarr.shape}, dtype={self._ndarr.dtype})'
 
-  @property
-  def ndarr(self) -> np.ndarray:
-    return self._ndarr
-
-  def setitem(self, target: np.ndarray) -> None:
-    """Get raw-data"""
-    print(
-        f'TODO: setitem, taget: {type(target)}, {target.shape}, {target.dtype}'
+  def _get_np_view(self, dataptr: int) -> np.ndarray:
+    """Create numpy array that uses `dataptr` memory block"""
+    ct_array = ctypes.cast(
+        dataptr, ctypes.POINTER(np.ctypeslib.as_ctypes_type(self._ndarr.dtype))
     )
-    target[...] = self._ndarr
+    return np.ctypeslib.as_array(ct_array, shape=self._ndarr.shape)
+
+  def setitem(self, src: 'SampleScalar', dataptr: int) -> None:
+    """Python `NPY_DT_setitem` implementation
+
+    NOTE:
+      `dataptr` is a buffer address, valid during execution of the function only
+    """
+    if not self.is_compatible(src):
+      raise ValueError('Incompatible item value')
+
+    self._get_np_view(dataptr)[...] = src._ndarr
+
+  def getitem(self, dataptr: int) -> 'SampleScalar':
+    """Python `NPY_DT_getitem` implementation
 
-  def getitem(self, source: np.ndarray) -> 'SampleScalar':
-    """Store raw-data into a new object"""
+    NOTE: See setitem()
+    """
     new = self.copy()
-    new._ndarr.flat[:] = source
+    new._ndarr[...] = self._get_np_view(dataptr)
     return new

src/dtype.c

@@ -103,37 +103,6 @@ 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,
@@ -145,26 +114,8 @@ static int sampledtype_setitem(SampleDTypeObject *descr, PyObject *obj,
     return -1;
   }
 
-  PyObject *res =
-      PyObject_CallMethod(descr->sample_scalar, "is_compatible", "O", obj);
-  if (res == NULL) {
-    return -1;
-  }
-  int is_compatible = PyObject_IsTrue(res);
-  Py_DECREF(res);
-
-  if (!is_compatible) {
-    PyErr_Format(PyExc_ValueError, "Incompatible item value");
-    return -1;
-  }
-
-  PyObject *view = create_element_view(descr, dataptr, NPY_FALSE);
-  if (view == NULL) {
-    return -1;
-  }
-
-  res = PyObject_CallMethod(obj, "setitem", "O", view);
-  Py_DECREF(view);
+  PyObject *res = PyObject_CallMethod(descr->sample_scalar, "setitem", "On",
+                                      obj, (Py_ssize_t)dataptr);
   if (res == NULL) {
     return -1;
   }
@@ -177,19 +128,8 @@ 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);
 
-  PyObject *view = create_element_view(descr, dataptr, NPY_TRUE);
-  if (view == NULL) {
-    return NULL;
-  }
-
-  PyObject *res =
-      PyObject_CallMethod(descr->sample_scalar, "getitem", "O", view);
-  Py_DECREF(view);
-  if (res == NULL) {
-    return NULL;
-  }
-
-  return res;
+  return PyObject_CallMethod(descr->sample_scalar, "getitem", "n",
+                             (Py_ssize_t)dataptr);
 }
 
 static PyType_Slot SampleDType_Slots[] = {

tests/test_create.py

@@ -28,7 +28,7 @@ def test_setitem():
   import sample_dtypes
 
   print('* create empty array')
-  scalar = sample_dtypes.SampleScalar(dtype=int)
+  scalar = sample_dtypes.SampleScalar(shape=(3,), dtype=int)
 
   print('  * set compatible value')
   arr: np.ndarray = np.empty(3, sample_dtypes.SampleDType(scalar))