Hi! I'm working on a Python C extension and I would like to expose a custom enum (as in: a class inheriting from enum.Enum) that would be entirely defined in C.
It turned out to not be a trivial task and the regular mechanism for inheritance using .tp_base doesn't work - most likely due to the Enum's meta class not being pulled in. Basically I'm trying to do this: import enum class FooBar(enum.Enum): FOO = 1 BAR = 2 in C. After a lot of digging into cpython's internals, this is what I came up with, wrapped in an example buildable module: #include <Python.h> PyDoc_STRVAR(module_doc, "C extension module defining a class inheriting from enum.Enum."); static PyModuleDef module_def = { PyModuleDef_HEAD_INIT, .m_name = "pycenum", .m_doc = module_doc, .m_size = -1, }; struct enum_descr { const char *name; long value; }; static const struct enum_descr foobar_descr[] = { { .name = "FOO", .value = 1, }, { .name = "BAR", .value = 2, }, { } }; static PyObject *make_bases(PyObject *enum_mod) { PyObject *enum_type, *bases; enum_type = PyObject_GetAttrString(enum_mod, "Enum"); if (!enum_type) return NULL; bases = PyTuple_Pack(1, enum_type); /* Steals reference. */ if (!bases) Py_DECREF(enum_type); return bases; } static PyObject *make_classdict(PyObject *enum_mod, PyObject *bases) { PyObject *enum_meta_type, *classdict; enum_meta_type = PyObject_GetAttrString(enum_mod, "EnumMeta"); if (!enum_meta_type) return NULL; classdict = PyObject_CallMethod(enum_meta_type, "__prepare__", "sO", "FooBarEnum", bases); Py_DECREF(enum_meta_type); return classdict; } static int fill_classdict(PyObject *classdict, PyObject *modname, const struct enum_descr *descr) { const struct enum_descr *entry; PyObject *key, *val; int ret; key = PyUnicode_FromString("__module__"); if (!key) return -1; ret = PyObject_SetItem(classdict, key, modname); Py_DECREF(key); if (ret < 0) return -1; for (entry = descr; entry->name; entry++) { key = PyUnicode_FromString(entry->name); if (!key) return -1; val = PyLong_FromLong(entry->value); if (!val) { Py_DECREF(key); return -1; } ret = PyObject_SetItem(classdict, key, val); Py_DECREF(key); Py_DECREF(val); if (ret < 0) return -1; } return 0; } static PyObject *make_new_type(PyObject *classdict, PyObject *bases, const char *enum_name) { PyObject *name, *args, *new_type; int ret; name = PyUnicode_FromString(enum_name); if (!name) return NULL; args = PyTuple_Pack(3, name, bases, classdict); if (!args) { Py_DECREF(name); return NULL; } Py_INCREF(bases); Py_INCREF(classdict); /* * Reference to name was stolen by PyTuple_Pack(), no need to * increase it here. */ new_type = PyObject_CallObject((PyObject *)&PyType_Type, args); Py_DECREF(args); if (!new_type) return NULL; ret = PyType_Ready((PyTypeObject *)new_type); if (ret < 0) { Py_DECREF(new_type); return NULL; } return new_type; } static PyObject *make_enum_type(PyObject *modname, const char *enum_name, const struct enum_descr *descr) { PyObject *enum_mod, *bases, *classdict, *new_type; int ret; enum_mod = PyImport_ImportModule("enum"); if (!enum_mod) return NULL; bases = make_bases(enum_mod); if (!bases) { Py_DECREF(enum_mod); return NULL; } classdict = make_classdict(enum_mod, bases); if (!classdict) { Py_DECREF(bases); Py_DECREF(enum_mod); return NULL; } ret = fill_classdict(classdict, modname, descr); if (ret < 0) { Py_DECREF(bases); Py_DECREF(enum_mod); Py_DECREF(classdict); return NULL; } new_type = make_new_type(classdict, bases, enum_name); Py_DECREF(bases); Py_DECREF(enum_mod); Py_DECREF(classdict); return new_type; } PyMODINIT_FUNC PyInit_pycenum(void) { PyObject *module, *modname, *sub_enum_type; int ret; module = PyModule_Create(&module_def); if (!module) return NULL; ret = PyModule_AddStringConstant(module, "__version__", "0.0.1"); if (ret < 0) { Py_DECREF(module); return NULL; } modname = PyModule_GetNameObject(module); if (!modname) { Py_DECREF(module); return NULL; } sub_enum_type = make_enum_type(modname, "FooBar", foobar_descr); Py_DECREF(modname); if (!sub_enum_type) { Py_DECREF(module); return NULL; } ret = PyModule_AddObject(module, "FooBar", sub_enum_type); if (ret < 0) { Py_DECREF(sub_enum_type); Py_DECREF(module); return NULL; } return module; } Basically I'm calling the EnumMeta's __prepare__ method directly to create a correct classdict and then I call the PyType_Type object too to create the sub-type. This works and AFAICT results in a class that behaves exactly as expected, but... am I doing this right? Any feedback is appreciated. I want to use this in a real project, namely the v2 of libgpiod python bindings[1] so it's important to me to get this right. Best regards, Bartosz Golaszewski [1] https://git.kernel.org/pub/scm/libs/libgpiod/libgpiod.git/ -- https://mail.python.org/mailman/listinfo/python-list