Estoy trabajando en una extensión de Python C y me gustaría exponer una enumeración personalizada (como en: una clase que hereda de enum.Enum) que estaría completamente definida en C.

Resultó no ser una tarea trivial y el mecanismo regular de herencia usando .tp_base no funciona, muy probablemente debido a que la metaclase de Enum no se incorporó.

Básicamente estoy tratando de hacer esto:

 import enum class FooBar(enum.Enum): FOO = 1 BAR = 2

C ª.

Después de indagar mucho en las partes internas de cpython, esto es lo que se me ocurrió, envuelto en un módulo construible de ejemplo:

 #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; }

Básicamente, estoy llamando al método EnumMeta de __prepare__ directamente para crear un classdict correcto y luego llamo al objeto PyType_Type también para crear el subtipo.

Esto funciona y AFAICT da como resultado una clase que se comporta exactamente como se esperaba, pero... ¿lo estoy haciendo bien? Cualquier comentario es apreciado.