Implement StringScalar in C

Author: ngoldbaum

stringdtype/meson.build

@@ -30,12 +30,13 @@ srcs = [
   'stringdtype/src/static_string.h',
   'stringdtype/src/umath.c',
   'stringdtype/src/umath.h',
+  'stringdtype/src/scalar.h',
+  'stringdtype/src/scalar.c',
 ]
 
 py.install_sources(
   [
     'stringdtype/__init__.py',
-    'stringdtype/scalar.py'
   ],
   subdir: 'stringdtype'
 )

stringdtype/pyproject.toml

@@ -31,6 +31,6 @@ per-file-ignores = {"__init__.py" = ["F401"]}
 
 [tool.meson-python.args]
 dist = []
-setup = ["-Ddebug=true", "-Doptimization=0"]
+setup = ["-Ddebug=true", "-Doptimization=2"]
 compile = []
 install = []

stringdtype/stringdtype/__init__.py

@@ -2,8 +2,7 @@
 
 """
 
-from .scalar import StringScalar  # isort: skip
-from ._main import StringDType, _memory_usage
+from ._main import StringDType, StringScalar, _memory_usage
 
 __all__ = [
     "StringDType",

stringdtype/stringdtype/scalar.py

@@ -1,10 +1,9 @@
 #include "dtype.h"
 
 #include "casts.h"
+#include "scalar.h"
 #include "static_string.h"
 
-PyTypeObject *StringScalar_Type = NULL;
-
 /*
  * Internal helper to create new instances
  */
@@ -64,7 +63,7 @@ static PyArray_Descr *
 string_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls),
                                          PyObject *obj)
 {
-    if (Py_TYPE(obj) != StringScalar_Type) {
+    if (Py_TYPE(obj) != &StringScalar_Type) {
         PyErr_SetString(PyExc_TypeError,
                         "Can only store StringScalar in a StringDType array.");
         return NULL;
@@ -84,7 +83,7 @@ get_value(PyObject *scalar)
     PyTypeObject *scalar_type = Py_TYPE(scalar);
     // FIXME: handle bytes too
     if ((scalar_type == &PyUnicode_Type) ||
-        (scalar_type == StringScalar_Type)) {
+        (scalar_type == &StringScalar_Type)) {
         // attempt to decode as UTF8
         ret_bytes = PyUnicode_AsUTF8String(scalar);
         if (ret_bytes == NULL) {
@@ -160,12 +159,7 @@ stringdtype_getitem(StringDTypeObject *NPY_UNUSED(descr), char **dataptr)
         return NULL;
     }
 
-    PyObject *res = PyObject_CallFunctionObjArgs((PyObject *)StringScalar_Type,
-                                                 val_obj, NULL);
-
-    Py_DECREF(val_obj);
-
-    return res;
+    return val_obj;
 }
 
 // PyArray_NonzeroFunc
@@ -400,7 +394,7 @@ init_string_dtype(void)
     PyArrayMethod_Spec **casts = get_casts();
 
     PyArrayDTypeMeta_Spec StringDType_DTypeSpec = {
-            .typeobj = StringScalar_Type,
+            .typeobj = &StringScalar_Type,
             .slots = StringDType_Slots,
             .casts = casts,
     };

stringdtype/stringdtype/src/dtype.c

@@ -18,7 +18,6 @@ typedef struct {
 } StringDTypeObject;
 
 extern PyArray_DTypeMeta StringDType;
-extern PyTypeObject *StringScalar_Type;
 
 StringDTypeObject *
 new_stringdtype_instance(void);

stringdtype/stringdtype/src/dtype.h

@@ -6,6 +6,7 @@
 #include "numpy/experimental_dtype_api.h"
 
 #include "dtype.h"
+#include "scalar.h"
 #include "static_string.h"
 #include "umath.h"
 
@@ -100,23 +101,24 @@ PyInit__main(void)
         return NULL;
     }
 
-    PyObject *mod = PyImport_ImportModule("stringdtype");
-    if (mod == NULL) {
+    if (init_stringdtype_scalar() < 0) {
         goto error;
     }
-    StringScalar_Type =
-            (PyTypeObject *)PyObject_GetAttrString(mod, "StringScalar");
-    Py_DECREF(mod);
 
-    if (StringScalar_Type == NULL) {
+    if (init_string_dtype() < 0) {
         goto error;
     }
 
-    if (init_string_dtype() < 0) {
+    Py_INCREF((PyObject *)&StringDType);
+    if (PyModule_AddObject(m, "StringDType", (PyObject *)&StringDType) < 0) {
+        Py_DECREF((PyObject *)&StringDType);
         goto error;
     }
 
-    if (PyModule_AddObject(m, "StringDType", (PyObject *)&StringDType) < 0) {
+    Py_INCREF((PyObject *)&StringScalar_Type);
+    if (PyModule_AddObject(m, "StringScalar", (PyObject *)&StringScalar_Type) <
+        0) {
+        Py_DECREF((PyObject *)&StringScalar_Type);
         goto error;
     }
 

stringdtype/stringdtype/src/main.c

@@ -0,0 +1,24 @@
+#define PY_ARRAY_UNIQUE_SYMBOL STRINGDTYPE_ARRAY_API
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#define NO_IMPORT_ARRAY
+#include "numpy/arrayobject.h"
+#include "numpy/experimental_dtype_api.h"
+#include "numpy/ndarraytypes.h"
+
+#include "scalar.h"
+
+PyTypeObject StringScalar_Type = {
+        PyVarObject_HEAD_INIT(NULL, 0).tp_name = "stringdtype.StringScalar",
+        .tp_basicsize = sizeof(StringScalarObject),
+};
+
+int
+init_stringdtype_scalar(void)
+{
+    StringScalar_Type.tp_base = &PyUnicode_Type;
+    if (PyType_Ready(&StringScalar_Type) < 0) {
+        return -1;
+    }
+
+    return 0;
+}

stringdtype/stringdtype/src/scalar.c

@@ -0,0 +1,15 @@
+#ifndef _STRINGDTYPE_SCALAR_H
+#define _STRINGDTYPE_SCALAR_H
+
+#include <Python.h>
+
+extern PyTypeObject StringScalar_Type;
+
+int
+init_stringdtype_scalar(void);
+
+typedef struct {
+    PyUnicodeObject str;
+} StringScalarObject;
+
+#endif /* _STRINGDTYPE_SCALAR_H */