Experimental support for dtype.hasobject

trundev · trundev · commit 7e2f8af7ff83 · 2025-09-14T13:19:10.000+03:00
Add NPY_DT_get_clear_loop implementation in python.
Include test for the key `custom_as_ctypes_type()` conversion.
diff --git a/sample_dtypes/scalar.py b/sample_dtypes/scalar.py
@@ -8,12 +8,44 @@
 BYTES_BUF_ORDER = 'A'
 
 
+def custom_as_ctypes_type(dtype: np.dtype):
+  """Wrapper around `numpy.ctypeslib.as_ctypes_type` that maps 'O' to py_object"""
+  if not dtype.hasobject:
+    return np.ctypeslib.as_ctypes_type(dtype)
+
+  if dtype.fields is None:
+    # Simple dtype or array
+    if dtype.subdtype:
+      # Handle sub-array types like ('O', 3)
+      base, shape = dtype.subdtype
+      new_base = custom_as_ctypes_type(base)
+      for dim in reversed(shape):
+        new_base = new_base * dim
+      return new_base
+    else:
+      # Simple object (non-objects are already handled)
+      assert dtype == np.dtype(object)
+      return ctypes.py_object
+
+  # Structured dtype
+  new_fields = []
+  for name, (field_dtype, *_) in dtype.fields.items():
+    new_field_dtype = custom_as_ctypes_type(field_dtype)
+    new_fields.append((name, new_field_dtype))
+
+  class CustomStruct(ctypes.Structure):
+    _pack_ = dtype.alignment
+    _fields_ = new_fields
+
+  return CustomStruct
+
+
 class SampleScalar:
   _ndarr: np.ndarray
 
   def __init__(self, *, shape: tuple[int, ...] = DEFAULT_SHAPE, dtype=None):
     print(f'TODO: scalar.py: SampleScalar.__init__({shape=}, {dtype=})')
-    self._ndarr = np.zeros(shape, dtype=dtype)
+    self._ndarr = np.empty(shape, dtype=dtype)
 
   def copy(self) -> 'SampleScalar':
     return type(self)(shape=self._ndarr.shape, dtype=self._ndarr.dtype)
@@ -26,6 +58,10 @@ def elsize(self) -> int:
   def alignment(self) -> int:
     return self._ndarr.dtype.alignment
 
+  @property
+  def hasobject(self) -> bool:
+    return self._ndarr.dtype.hasobject
+
   def is_compatible(self, other: 'SampleScalar') -> bool:
     print(f'scalar.py: SampleScalar.is_compatible({other})')
     return self._ndarr.shape == other._ndarr.shape
@@ -36,10 +72,35 @@ def __repr__(self) -> str:
   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))
+        dataptr, ctypes.POINTER(custom_as_ctypes_type(self._ndarr.dtype))
     )
     return np.ctypeslib.as_array(ct_array, shape=self._ndarr.shape)
 
+  def _get_strided_np_view(
+      self, dataptr: int, size: int, stride: int
+  ) -> np.ndarray:
+    """Create numpy array that uses strided `dataptr` memory block"""
+    # Start with ctypes object to cover single SampleScalar
+    ct_base = custom_as_ctypes_type(self._ndarr.dtype)
+    for dim in reversed(self._ndarr.shape):
+      ct_base = ct_base * dim
+
+    # Single ctypes object to cover whole strided data-block
+    # CHECKME: `stride` may also be zero (?), but `size` must be 1
+    assert stride >= ctypes.sizeof(ct_base), (
+        'Insufficient stride size ' f'{stride} / {ctypes.sizeof(ct_base)}'
+    )
+
+    class PaddedStruct(ctypes.Structure):
+      _align_ = stride
+      _fields_ = [('_', ct_base)]
+
+    ct_base = PaddedStruct * size
+
+    # Create ndarray view from this ctypes object
+    ct_array = ctypes.cast(dataptr, ctypes.POINTER(ct_base))
+    return np.ctypeslib.as_array(ct_array, shape=())
+
   def setitem(self, src: 'SampleScalar', dataptr: int) -> None:
     """Python `NPY_DT_setitem` implementation
 
@@ -59,3 +120,13 @@ def getitem(self, dataptr: int) -> 'SampleScalar':
     new = self.copy()
     new._ndarr[...] = self._get_np_view(dataptr)
     return new
+
+  def clear_loop(self, data: int, size: int, stride: int) -> int:
+    """Python implementation for `out_loop` returned by `NPY_DT_get_clear_loop`
+
+    NOTE: See setitem()
+    """
+    print(f'scalar.py: SampleScalar.clear_loop {data=}, {size=}, {stride=}')
+    view = self._get_strided_np_view(data, size, stride)
+    view[...] = 0  # Force dereference of all object-entries
+    return 0
diff --git a/src/dtype.c b/src/dtype.c
@@ -39,6 +39,16 @@ SampleDTypeObject *new_sampledtype_instance(PyObject *sample_scalar) {
     return NULL;
   }
 
+  attr = PyObject_GetAttrString(sample_scalar, "hasobject");
+  if (attr == NULL) {
+    return NULL;
+  }
+  int hasobject = PyObject_IsTrue(attr);
+  Py_DECREF(attr);
+  if (hasobject == -1 && PyErr_Occurred()) {
+    return NULL;
+  }
+
   // Allocate final DType bject
   SampleDTypeObject *new = (SampleDTypeObject *)PyArrayDescr_Type.tp_new(
       (PyTypeObject *)&SampleDType, NULL, NULL);
@@ -48,6 +58,10 @@ SampleDTypeObject *new_sampledtype_instance(PyObject *sample_scalar) {
 
   Py_INCREF(sample_scalar);
   new->sample_scalar = sample_scalar;
+  if (hasobject) {
+    // TODO: Other flags from NPY_OBJECT_DTYPE_FLAGS
+    new->base.flags |= NPY_ITEM_HASOBJECT | NPY_NEEDS_INIT;
+  }
   new->base.elsize = elsize;
   new->base.alignment = alignment;
 
@@ -105,7 +119,7 @@ sampledtype_ensure_canonical(SampleDTypeObject *self) {
 
 static int sampledtype_setitem(SampleDTypeObject *descr, PyObject *obj,
                                char *dataptr) {
-  printf("%s, target elsise %lld, type_num %d\n", __func__, descr->base.elsize,
+  printf("%s, target elsise %zd, type_num %d\n", __func__, descr->base.elsize,
          descr->base.type_num);
 
   if (Py_TYPE(obj) != SampleScalar_Type) {
@@ -125,13 +139,51 @@ static int sampledtype_setitem(SampleDTypeObject *descr, PyObject *obj,
 }
 
 static PyObject *sampledtype_getitem(SampleDTypeObject *descr, char *dataptr) {
-  printf("%s, source elsize %lld, type_num %d\n", __func__, descr->base.elsize,
+  printf("%s, source elsize %zd, type_num %d\n", __func__, descr->base.elsize,
          descr->base.type_num);
 
   return PyObject_CallMethod(descr->sample_scalar, "getitem", "n",
                              (Py_ssize_t)dataptr);
 }
 
+static int sampledtype_clear_loop(void *NPY_UNUSED(traverse_context),
+                                  const PyArray_Descr *_descr, char *data,
+                                  npy_intp size, npy_intp stride,
+                                  NpyAuxData *NPY_UNUSED(auxdata)) {
+  const SampleDTypeObject *descr = (const SampleDTypeObject *)_descr;
+  printf("%s, source elsize %zd, type_num %d, data %p, size %zd, stride %zd\n",
+         __func__, descr->base.elsize, descr->base.type_num, data, size,
+         stride);
+
+  PyObject *res = PyObject_CallMethod(descr->sample_scalar, "clear_loop", "nnn",
+                                      (Py_ssize_t)data, size, stride);
+  if (res == NULL) {
+    return -1;
+  }
+  int ret = PyLong_AsLong(res);
+  Py_DECREF(res);
+  if (ret == -1 && PyErr_Occurred()) {
+    return -1;
+  }
+
+  return ret;
+}
+
+static int sampledtype_get_clear_loop(void *NPY_UNUSED(traverse_context),
+                                      const SampleDTypeObject *descr,
+                                      int aligned, npy_intp fixed_stride,
+                                      PyArrayMethod_TraverseLoop **out_loop,
+                                      NpyAuxData **NPY_UNUSED(out_auxdata),
+                                      NPY_ARRAYMETHOD_FLAGS *flags) {
+  printf("%s, source elsize %zd, type_num %d, aligned %d, fixed_stride %zd\n",
+         __func__, descr->base.elsize, descr->base.type_num, aligned,
+         fixed_stride);
+
+  *out_loop = &sampledtype_clear_loop;
+  *flags = NPY_METH_NO_FLOATINGPOINT_ERRORS;
+  return 0;
+}
+
 static PyType_Slot SampleDType_Slots[] = {
     {NPY_DT_discover_descr_from_pyobject,
      &sampledtype_discover_descriptor_from_pyobject},
@@ -140,6 +192,7 @@ static PyType_Slot SampleDType_Slots[] = {
     {NPY_DT_ensure_canonical, &sampledtype_ensure_canonical},
     {NPY_DT_setitem, &sampledtype_setitem},
     {NPY_DT_getitem, &sampledtype_getitem},
+    {NPY_DT_get_clear_loop, &sampledtype_get_clear_loop},
     {0, NULL}};
 
 static PyObject *sampledtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args,
diff --git a/tests/test_refcount.py b/tests/test_refcount.py
@@ -1,16 +1,18 @@
-"""Test internal object reference counting"""
+"""Test internal object reference counting
+
+NOTE:
+  Most of these tests fail when debugging, because "ms-python.debugpy" takes
+  extra local object ownership.
+"""
 
 import sys
 import numpy as np
+import pytest
+import sample_dtypes
 
 
 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
-
+  """Test SampleScalar internal dtype reference counts"""
   # Use non-singleton (non-immortal) dtypes to have valid ref-counts
   dtype = np.dtype([('a', np.int16)])
   assert (
@@ -23,3 +25,80 @@ def test_dtype_refcnt():
 
   del arr, scalar
   assert sys.getrefcount(dtype) == 2, 'Unexpected dtype ref-count'
+
+
+@pytest.mark.parametrize(
+    'dtype',
+    [
+        # Simple object dtype
+        np.dtype(object),
+        # Array of objects
+        np.dtype((object, 3)),
+        # Structured dtype to include object (alignment issue)
+        np.dtype([('i', np.int16), ('obj', object)]),
+        # Complex/compound dtype with Structured and Array elements
+        np.dtype(
+            ([('i', np.uint16), ('obj_arr', [('obj', object)], (2, 3))], 4)
+        ),
+    ],
+)
+def test_custom_as_ctypes_type(dtype):
+  """Test scalar.custom_as_ctypes_type conversion, see np.ctypeslib.as_ctypes_type"""
+  import ctypes
+  from sample_dtypes.scalar import custom_as_ctypes_type
+
+  ct_type = custom_as_ctypes_type(dtype)
+  assert (
+      ctypes.sizeof(ct_type) == dtype.itemsize
+  ), 'Converted size does NOT match'
+
+  shape = (3, 2)
+  src = np.empty(shape, dtype=dtype)
+  tgt = np.empty(shape, dtype=dtype)
+  # Generate some different source values
+  src.flat[...] = (np.arange(src.size) + 10).astype(src.dtype)
+
+  # Create view that uses `tgt` data-buffer
+  ct_array = ctypes.cast(
+      tgt.__array_interface__['data'][0], ctypes.POINTER(ct_type)
+  )
+  tgt_view = np.ctypeslib.as_array(ct_array, shape=shape)
+  # Replace view content, expect same content in `tgt`
+  tgt_view[...] = src
+  np.testing.assert_array_equal(tgt, src, 'Mismatch between array and its view')
+
+
+def test_sub_dtype_refcnt():
+  """Test SampleScalar internal dtype reference counts"""
+  # Use generic-object dtype to test sub-object ownership transfer
+  scalar = sample_dtypes.SampleScalar(dtype=object)
+  arr: np.ndarray = np.empty(3, sample_dtypes.SampleDType(scalar.copy()))
+
+  obj = f'string {1}'  # Dynamic (non-immortal) string object
+  scalar._ndarr[0, 1] = obj
+  # Expect 3 references to the string: `obj`, `scalar._ndarr` and `getrefcount`
+  assert (
+      sys.getrefcount(obj) == 3
+  ), 'Unexpected ref-count for non-immortal object'
+
+  # Do NPY_DT_setitem
+  arr[1] = scalar
+  # Expect 4 references to the string: ones from above and `arr._ndarr`
+  assert sys.getrefcount(obj) == 4, 'Unexpected ref-count after assignment'
+  assert obj is arr[1]._ndarr[0, 1], 'Same object must be referenced by arr'
+  arr[2] = scalar
+  # Expect 5 references to the string: ones from above and extra `arr._ndarr`
+  assert sys.getrefcount(obj) == 5, 'Unexpected ref-count after assignment'
+  assert obj is arr[2]._ndarr[0, 1], 'Same object must be referenced by arr'
+
+  # Drop references one by one
+  del scalar
+  # Expect `scalar._ndarr` reference to be dropped
+  assert (
+      sys.getrefcount(obj) == 4
+  ), 'Unexpected ref-count after scalar reference deletion'
+  del arr
+  # Expect two `arr._ndarr` references to be dropped
+  assert (
+      sys.getrefcount(obj) == 2
+  ), 'Unexpected ref-count after ndarray reference deletion'
