Experiment: add register_dtype and some_dtype.h

Author: trundev

sample_dtypes/__init__.py

@@ -3,4 +3,17 @@
 __version__ = "0.0.0"  # Keep in sync with pyproject.toml:version
 __all__ = [
     "__version__",
+    "register_dtype",
 ]
+
+import sample_dtypes
+from sample_dtypes._sample_dtypes_ext import _register_dtype_ext
+
+
+def register_dtype(name: str, data_type: type, count: int):
+  dtype = _register_dtype_ext(name, data_type, count)
+  setattr(sample_dtypes, name, dtype)
+  setattr(dtype, "__module__", sample_dtypes)
+  # Support dtype(type_name)
+  # setattr(dtype, 'dtype', npy_descr)
+  return dtype

setup.py

@@ -8,11 +8,11 @@
 
 if platform.system() == "Windows":
   COMPILE_ARGS = [
-      "/std:c++17",
+      "/std:c++20",
   ]
 else:
   COMPILE_ARGS = [
-      "-std=c++17",
+      "-std=c++20",
   ]
 
 

src/dtypes.cc

@@ -8,9 +8,84 @@
 #include <Python.h>
 
 #include "src/common.h"
+#include "src/some_dtype.h"
 
 namespace sample_dtypes {
 
+static PyObject *RegisterDtype(PyObject *self, PyObject *args) {
+  std::cout << "DEBUG: " << __func__ << std::endl;
+
+  const char *name = nullptr;
+  unsigned int count = 0;
+  PyTypeObject *data_type = nullptr;
+  if (!PyArg_ParseTuple(args, "sO!I", &name, &PyType_Type, &data_type,
+                        &count)) {
+    std::cout << "DEBUG: " << "PyArg_ParseTuple failed" << std::endl;
+    return nullptr;
+  }
+
+  // Get the size of type
+  unsigned data_type_size = 0;
+  if (PyType_IsSubtype(data_type, &PyLong_Type)) {
+    data_type_size = sizeof(uint64_t);
+    std::cout << "DEBUG: data_type(int) specified" << std::endl;
+  } else if (PyType_IsSubtype(data_type, &PyFloat_Type)) {
+    data_type_size = sizeof(float);
+    std::cout << "DEBUG: data_type(float) specified" << std::endl;
+  } else {
+    // Not a supported python data-type, check numpy types
+    // PyObject *type_dict = PyType_GetDict(type);
+    PyObject *type_nbytes = PyObject_GetAttrString(
+        reinterpret_cast<PyObject *>(data_type), "nbytes");
+    if (type_nbytes == nullptr) {
+      std::cout << "DEBUG: PyObject_GetAttrString 'nbytes' failed" << std::endl;
+      PyErr_Format(PyExc_TypeError, "Cannot interpret '%R' as a data type",
+                   data_type);
+      return nullptr;
+    }
+    if (!PyArg_Parse(type_nbytes, "I", &data_type_size)) {
+      std::cout << "DEBUG: PyArg_Parse failed" << std::endl;
+      return nullptr;
+    }
+    std::cout << "DEBUG: numpy.dtype specified, nbytes=" << data_type_size
+              << std::endl;
+  }
+
+  PyObject *type_name = PyType_GetQualName(data_type);
+  const char *type_name_str = nullptr;
+  if (!PyArg_Parse(type_name, "s", &type_name_str)) {
+    std::cout << "DEBUG: PyArg_Parse failed" << std::endl;
+    return nullptr;
+  }
+  std::cout << "DEBUG: _register_dtype_ext(name=" << name
+            << ", data_type=" << type_name_str << ", count=" << count
+            << ") -> data_type_size=" << data_type_size << std::endl;
+
+  // Create python type object
+  PyType_Slot type_slots[]{
+      {0, nullptr},
+  };
+  typedef struct {
+    PyObject_HEAD
+    // int value[0];
+  } type_struct;
+  PyType_Spec type_spec = {
+      .name = name,
+      .basicsize =
+          static_cast<int>(sizeof(type_struct) + data_type_size * count),
+      .itemsize = 0,
+      .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+      .slots = type_slots,
+  };
+  PyObject *type = PyType_FromSpecWithBases(
+      &type_spec, reinterpret_cast<PyObject *>(&PyGenericArrType_Type));
+  if (!type) {
+    return nullptr;
+  }
+
+  return type;
+}
+
 /*
  *TODO: Implement Dtype casts registration
  */
@@ -49,13 +124,25 @@ bool Initialize() {
   return success;
 }
 
+static PyMethodDef module_methods[] = {{"_register_dtype_ext", RegisterDtype,
+                                        METH_VARARGS,
+                                        "Create/register a dtype class"},
+                                       {nullptr, nullptr, 0}};
+
 static PyModuleDef module_def = {
     PyModuleDef_HEAD_INIT,
     "_sample_dtypes_ext",
+    "sample_dtypes extension module",
+    -1,
+    module_methods,
 };
 
 PyMODINIT_FUNC PyInit__sample_dtypes_ext() {
-  std::cout << "DEBUG: " << __func__ << std::endl;
+  std::cout << "DEBUG: " << __DATE__ << " " << __TIME__ << ": " << __func__
+            << std::endl;
+#ifdef _HYBRID
+  std::cout << "HYBRID build" << std::endl;
+#endif
 
   Safe_PyObjectPtr m = make_safe(PyModule_Create(&module_def));
   if (!m) {

src/some_dtype.h

@@ -0,0 +1,12 @@
+#ifndef SAMPLE_DTYPES_SOME_DTYPE_H_
+#define SAMPLE_DTYPES_SOME_DTYPE_H_
+
+namespace sample_dtypes {
+
+template <int N, typename T> struct some_dtype {
+  T v[N];
+};
+
+} // namespace sample_dtypes
+
+#endif // SAMPLE_DTYPES_SOME_DTYPE_H_

tests/test_register.py

@@ -0,0 +1,14 @@
+"""Test register_dtype()"""
+
+
+def test_register_dtype():
+  from sample_dtypes import _sample_dtypes_ext as pyd
+
+  print(
+      f'_register_dtype_ext: {type(pyd._register_dtype_ext)}, {type(pyd.__doc__)}'
+  )
+
+  import sample_dtypes
+
+  dtype = sample_dtypes.register_dtype('type_name', int, 4)
+  assert dtype.__module__ is sample_dtypes