From d6310c740a067796743e97f9254c6d81f1f8e7d1 Mon Sep 17 00:00:00 2001 From: Nathan Goldbaum Date: Wed, 24 May 2023 12:12:20 -0600 Subject: [PATCH] WIP: convert MetadataDType to a heap type --- metadatadtype/metadatadtype/src/dtype.c | 137 ++++++++++++------ metadatadtype/metadatadtype/src/dtype.h | 18 ++- .../metadatadtype/src/metadatadtype_main.c | 76 +++++++--- metadatadtype/metadatadtype/src/umath.c | 17 +-- metadatadtype/metadatadtype/src/umath.h | 10 +- 5 files changed, 171 insertions(+), 87 deletions(-) diff --git a/metadatadtype/metadatadtype/src/dtype.c b/metadatadtype/metadatadtype/src/dtype.c index 9d4c6399..7398c4bc 100644 --- a/metadatadtype/metadatadtype/src/dtype.c +++ b/metadatadtype/metadatadtype/src/dtype.c @@ -2,14 +2,12 @@ #include "casts.h" -PyTypeObject *MetadataScalar_Type = NULL; - /* * `get_value` and `get_unit` are small helpers to deal with the scalar. */ static double -get_value(PyObject *scalar) +get_value(PyObject *scalar, const PyTypeObject *MetadataScalar_Type) { PyTypeObject *scalar_type = Py_TYPE(scalar); if (scalar_type != MetadataScalar_Type) { @@ -36,7 +34,7 @@ get_value(PyObject *scalar) } static PyObject * -get_metadata(PyObject *scalar) +get_metadata(PyObject *scalar, const PyTypeObject *MetadataScalar_Type) { if (Py_TYPE(scalar) != MetadataScalar_Type) { PyErr_SetString( @@ -64,11 +62,13 @@ get_metadata(PyObject *scalar) * Internal helper to create new instances */ MetadataDTypeObject * -new_metadatadtype_instance(PyObject *metadata) +new_metadatadtype_instance(metadatadtype_state *state, PyObject *metadata) { + if (state == NULL) { + return NULL; + } MetadataDTypeObject *new = (MetadataDTypeObject *)PyArrayDescr_Type.tp_new( - /* TODO: Using NULL for args here works, but seems not clean? */ - (PyTypeObject *)&MetadataDType, NULL, NULL); + (PyTypeObject *)state->MetadataDType, NULL, NULL); if (new == NULL) { return NULL; } @@ -111,17 +111,18 @@ common_dtype(PyArray_DTypeMeta *cls, PyArray_DTypeMeta *other) } static PyArray_Descr * -metadata_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls), +metadata_discover_descriptor_from_pyobject(PyArray_DTypeMeta *cls, PyObject *obj) { - if (Py_TYPE(obj) != MetadataScalar_Type) { + metadatadtype_state *state = PyType_GetModuleState((PyTypeObject *)cls); + if (Py_TYPE(obj) != state->MetadataScalar_Type) { PyErr_SetString( PyExc_TypeError, "Can only store MetadataScalar in a MetadataDType array."); return NULL; } - PyObject *metadata = get_metadata(obj); + PyObject *metadata = get_metadata(obj, state->MetadataScalar_Type); if (metadata == NULL) { return NULL; } @@ -135,7 +136,8 @@ metadata_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls), static int metadatadtype_setitem(MetadataDTypeObject *descr, PyObject *obj, char *dataptr) { - double value = get_value(obj); + metadatadtype_state *state = PyType_GetModuleState(Py_TYPE(descr)); + double value = get_value(obj, state->MetadataScalar_Type); if (value == -1 && PyErr_Occurred()) { return -1; } @@ -148,6 +150,7 @@ metadatadtype_setitem(MetadataDTypeObject *descr, PyObject *obj, char *dataptr) static PyObject * metadatadtype_getitem(MetadataDTypeObject *descr, char *dataptr) { + metadatadtype_state *state = PyType_GetModuleState(Py_TYPE(descr)); double val; /* get the value */ memcpy(&val, dataptr, sizeof(double)); // NOLINT @@ -158,7 +161,7 @@ metadatadtype_getitem(MetadataDTypeObject *descr, char *dataptr) } PyObject *res = PyObject_CallFunctionObjArgs( - (PyObject *)MetadataScalar_Type, val_obj, descr, NULL); + (PyObject *)state->MetadataScalar_Type, val_obj, descr, NULL); if (res == NULL) { return NULL; } @@ -186,9 +189,10 @@ static PyType_Slot MetadataDType_Slots[] = { {0, NULL}}; static PyObject * -metadatadtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, - PyObject *kwds) +metadatadtype_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { + metadatadtype_state *state = PyType_GetModuleState(type); + static char *kwargs_strs[] = {"metadata", NULL}; PyObject *metadata = NULL; @@ -201,13 +205,41 @@ metadatadtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, metadata = Py_None; } - return (PyObject *)new_metadatadtype_instance(metadata); + return (PyObject *)new_metadatadtype_instance(state, metadata); +} + +/* MetadataDType finalization */ + +static int +metadatadtype_traverse(PyObject *self_obj, visitproc visit, void *arg) +{ + // Visit the type + Py_VISIT(Py_TYPE(self_obj)); + + // Visit the metadata attribute + Py_VISIT(((MetadataDTypeObject *)self_obj)->metadata); + + return 0; +} + +static int +metadatadtype_clear(MetadataDTypeObject *self) +{ + Py_CLEAR(self->metadata); + return 0; +} + +static void +metadatadtype_finalize(PyObject *self_obj) +{ + Py_CLEAR(((MetadataDTypeObject *)self_obj)->metadata); } static void metadatadtype_dealloc(MetadataDTypeObject *self) { - Py_CLEAR(self->metadata); + PyObject_GC_UnTrack(self); + metadatadtype_finalize((PyObject *)self); PyArrayDescr_Type.tp_dealloc((PyObject *)self); } @@ -221,32 +253,52 @@ metadatadtype_repr(MetadataDTypeObject *self) static PyMemberDef MetadataDType_members[] = { {"_metadata", T_OBJECT_EX, offsetof(MetadataDTypeObject, metadata), READONLY, "some metadata"}, - {NULL}, + {NULL, 0, 0, 0, ""}, }; -/* - * This is the basic things that you need to create a Python Type/Class in C. - * However, there is a slight difference here because we create a - * PyArray_DTypeMeta, which is a larger struct than a typical type. - * (This should get a bit nicer eventually with Python >3.11.) - */ -PyArray_DTypeMeta MetadataDType = { - {{ - PyVarObject_HEAD_INIT(NULL, 0).tp_name = - "metadatadtype.MetadataDType", - .tp_basicsize = sizeof(MetadataDTypeObject), - .tp_new = metadatadtype_new, - .tp_dealloc = (destructor)metadatadtype_dealloc, - .tp_repr = (reprfunc)metadatadtype_repr, - .tp_str = (reprfunc)metadatadtype_repr, - .tp_members = MetadataDType_members, - }}, - /* rest, filled in during DTypeMeta initialization */ +static PyType_Slot MetadataDType_Type_slots[] = { + {Py_tp_traverse, metadatadtype_traverse}, + {Py_tp_clear, metadatadtype_clear}, + {Py_tp_finalize, metadatadtype_finalize}, + {Py_tp_dealloc, metadatadtype_dealloc}, + {Py_tp_str, (reprfunc)metadatadtype_repr}, + {Py_tp_repr, (reprfunc)metadatadtype_repr}, + {Py_tp_members, MetadataDType_members}, + {Py_tp_new, metadatadtype_new}, + {0, 0}, /* sentinel */ +}; + +static PyType_Spec MetadataDType_Type_spec = { + .name = "metadatadtype.MetadataDType", + .basicsize = sizeof(MetadataDTypeObject), + .flags = Py_TPFLAGS_DEFAULT, + .slots = MetadataDType_Type_slots, }; int -init_metadata_dtype(void) +init_metadata_dtype(PyObject *m) { + metadatadtype_state *state = PyModule_GetState(m); + + PyObject *bases = PyTuple_Pack(1, (PyObject *)&PyArrayDescr_Type); + + state->MetadataDType = (PyArray_DTypeMeta *)PyType_FromModuleAndSpec( + m, &MetadataDType_Type_spec, bases); + if (state->MetadataDType == NULL) { + return -1; + } + + // manually set type so it has the correct metaclass + ((PyObject *)state->MetadataDType)->ob_type = &PyArrayDTypeMeta_Type; + + // manually null remaining fields, is there a more forward compatible + // way to do this with e.g. memset? + state->MetadataDType->singleton = NULL; + state->MetadataDType->type_num = 0; + state->MetadataDType->scalar_type = NULL; + state->MetadataDType->flags = 0; + state->MetadataDType->dt_slots = 0; + /* * To create our DType, we have to use a "Spec" that tells NumPy how to * do it. You first have to create a static type, but see the note there! @@ -256,22 +308,17 @@ init_metadata_dtype(void) PyArrayDTypeMeta_Spec MetadataDType_DTypeSpec = { .flags = NPY_DT_PARAMETRIC | NPY_DT_NUMERIC, .casts = casts, - .typeobj = MetadataScalar_Type, + .typeobj = state->MetadataScalar_Type, .slots = MetadataDType_Slots, }; - /* Loaded dynamically, so may need to be set here: */ - ((PyObject *)&MetadataDType)->ob_type = &PyArrayDTypeMeta_Type; - ((PyTypeObject *)&MetadataDType)->tp_base = &PyArrayDescr_Type; - if (PyType_Ready((PyTypeObject *)&MetadataDType) < 0) { - return -1; - } - if (PyArrayInitDTypeMeta_FromSpec(&MetadataDType, + if (PyArrayInitDTypeMeta_FromSpec(state->MetadataDType, &MetadataDType_DTypeSpec) < 0) { return -1; } - MetadataDType.singleton = PyArray_GetDefaultDescr(&MetadataDType); + state->MetadataDType->singleton = + PyArray_GetDefaultDescr(state->MetadataDType); free(MetadataDType_DTypeSpec.casts[1]->dtypes); free(MetadataDType_DTypeSpec.casts[1]); diff --git a/metadatadtype/metadatadtype/src/dtype.h b/metadatadtype/metadatadtype/src/dtype.h index b3011a34..6436e098 100644 --- a/metadatadtype/metadatadtype/src/dtype.h +++ b/metadatadtype/metadatadtype/src/dtype.h @@ -13,24 +13,32 @@ #include "numpy/experimental_dtype_api.h" #include "numpy/ndarraytypes.h" +// module state +typedef struct { + PyArray_DTypeMeta *MetadataDType; + PyTypeObject *MetadataScalar_Type; +} metadatadtype_state; + +// MetadataDType objects + +// Instance state typedef struct { PyArray_Descr base; PyObject *metadata; } MetadataDTypeObject; -extern PyArray_DTypeMeta MetadataDType; -extern PyTypeObject *MetadataScalar_Type; - MetadataDTypeObject * -new_metadatadtype_instance(PyObject *metadata); +new_metadatadtype_instance(metadatadtype_state *state, PyObject *metadata); int -init_metadata_dtype(void); +init_metadata_dtype(PyObject *m); PyArray_Descr * common_instance(MetadataDTypeObject *dtype1, MetadataDTypeObject *NPY_UNUSED(dtype2)); +extern struct PyModuleDef metadatadtype_module; + // from numpy's dtypemeta.h, not publicly available #define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr)) diff --git a/metadatadtype/metadatadtype/src/metadatadtype_main.c b/metadatadtype/metadatadtype/src/metadatadtype_main.c index e7783fde..5ea0f4a8 100644 --- a/metadatadtype/metadatadtype/src/metadatadtype_main.c +++ b/metadatadtype/metadatadtype/src/metadatadtype_main.c @@ -8,56 +8,84 @@ #include "dtype.h" #include "umath.h" -static struct PyModuleDef moduledef = { - PyModuleDef_HEAD_INIT, - .m_name = "metadatadtype_main", - .m_size = -1, -}; - -/* Module initialization function */ -PyMODINIT_FUNC -PyInit__metadatadtype_main(void) +static int +metadatadtypemodule_modexec(PyObject *m) { + metadatadtype_state *state = PyModule_GetState(m); + if (_import_array() < 0) { - return NULL; + return -1; } if (import_experimental_dtype_api(11) < 0) { - return NULL; - } - - PyObject *m = PyModule_Create(&moduledef); - if (m == NULL) { - return NULL; + return -1; } PyObject *mod = PyImport_ImportModule("metadatadtype"); if (mod == NULL) { goto error; } - MetadataScalar_Type = + state->MetadataScalar_Type = (PyTypeObject *)PyObject_GetAttrString(mod, "MetadataScalar"); Py_DECREF(mod); - if (MetadataScalar_Type == NULL) { + if (state->MetadataScalar_Type == NULL) { goto error; } - if (init_metadata_dtype() < 0) { + if (init_metadata_dtype(m) < 0) { goto error; } - if (PyModule_AddObject(m, "MetadataDType", (PyObject *)&MetadataDType) < - 0) { + if (PyModule_AddType(m, (PyTypeObject *)state->MetadataDType) < 0) { goto error; } - if (init_ufuncs() < 0) { + if (init_ufuncs(state->MetadataDType) < 0) { goto error; } - return m; + return 0; error: Py_DECREF(m); - return NULL; + return -1; +} + +static int +metadatadtypemodule_traverse(PyObject *module, visitproc visit, void *arg) +{ + metadatadtype_state *state = PyModule_GetState(module); + Py_VISIT(state->MetadataDType); + Py_VISIT(state->MetadataScalar_Type); + return 0; +} + +static int +metadatadtypemodule_clear(PyObject *module) +{ + metadatadtype_state *state = PyModule_GetState(module); + Py_CLEAR(state->MetadataDType); + Py_CLEAR(state->MetadataScalar_Type); + return 0; +} + +static PyModuleDef_Slot metadatadtypemodule_slots[] = { + {Py_mod_exec, metadatadtypemodule_modexec}, + {0, NULL}, +}; + +struct PyModuleDef metadatadtype_module = { + PyModuleDef_HEAD_INIT, + .m_name = "_metadatadtype_main", + .m_size = sizeof(metadatadtype_state), + .m_slots = metadatadtypemodule_slots, + .m_traverse = metadatadtypemodule_traverse, + .m_clear = metadatadtypemodule_clear, +}; + +/* Module initialization function */ +PyMODINIT_FUNC +PyInit__metadatadtype_main(void) +{ + return PyModuleDef_Init(&metadatadtype_module); } diff --git a/metadatadtype/metadatadtype/src/umath.c b/metadatadtype/metadatadtype/src/umath.c index 1be362a5..ac665b2a 100644 --- a/metadatadtype/metadatadtype/src/umath.c +++ b/metadatadtype/metadatadtype/src/umath.c @@ -1,15 +1,8 @@ #include -#define PY_ARRAY_UNIQUE_SYMBOL metadatadtype_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 "numpy/ufuncobject.h" +#include "umath.h" #include "dtype.h" -#include "umath.h" static int translate_given_descrs(int nin, int nout, @@ -103,10 +96,10 @@ add_wrapping_loop(const char *ufunc_name, PyArray_DTypeMeta **dtypes, } int -init_ufuncs(void) +init_ufuncs(PyArray_DTypeMeta *MetadataDType) { - PyArray_DTypeMeta *binary_orig_dtypes[3] = {&MetadataDType, &MetadataDType, - &MetadataDType}; + PyArray_DTypeMeta *binary_orig_dtypes[3] = {MetadataDType, MetadataDType, + MetadataDType}; PyArray_DTypeMeta *binary_wrapped_dtypes[3] = { &PyArray_DoubleDType, &PyArray_DoubleDType, &PyArray_DoubleDType}; if (add_wrapping_loop("multiply", binary_orig_dtypes, @@ -114,7 +107,7 @@ init_ufuncs(void) goto error; } - PyArray_DTypeMeta *unary_boolean_dtypes[2] = {&MetadataDType, + PyArray_DTypeMeta *unary_boolean_dtypes[2] = {MetadataDType, &PyArray_BoolDType}; PyArray_DTypeMeta *unary_boolean_wrapped_dtypes[2] = {&PyArray_DoubleDType, &PyArray_BoolDType}; diff --git a/metadatadtype/metadatadtype/src/umath.h b/metadatadtype/metadatadtype/src/umath.h index 0a7e5e92..f57ad253 100644 --- a/metadatadtype/metadatadtype/src/umath.h +++ b/metadatadtype/metadatadtype/src/umath.h @@ -1,7 +1,15 @@ #ifndef _NPY_UFUNC_H #define _NPY_UFUNC_H +#define PY_ARRAY_UNIQUE_SYMBOL metadatadtype_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 "numpy/ufuncobject.h" + int -init_ufuncs(void); +init_ufuncs(PyArray_DTypeMeta *MetadataDType); #endif /*_NPY_UFUNC_H */