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.
The metaclass in Enum is tricky yes.
But you can see here that you can create an enum (in Python) like:
FooBar = enum.Enum('FooBar', dict(FOO=1, BAR=2))
So you can use this technique to easily create an enum class in Python C-API by doing something like:
PyObject *key, *val, *name, *attrs, *args, *modname, *kwargs, *enum_type, *sub_enum_type;
attrs = PyDict_New();
key = PyUnicode_FromString("FOO");
val = PyLong_FromLong(1);
PyObject_SetItem(attrs, key, val);
Py_DECREF(key);
Py_DECREF(val);
key = PyUnicode_FromString("BAR");
val = PyLong_FromLong(2);
PyObject_SetItem(attrs, key, val);
Py_DECREF(key);
Py_DECREF(val);
name = PyUnicode_FromString("FooBar");
args = PyTuple_Pack(3, name, attrs);
Py_DECREF(attrs);
Py_DECREF(name);
// the module name might need to be passed as keyword argument
PyDict_Type *kwargs = PyDict_New();
key = PyUnicode_FromString("module");
modname = PyModule_GetNameObject(module);
PyObject_SetItem(kwargs, key, modname);
Py_DECREF(key);
Py_DECREF(modname);
enum_type = PyObject_GetAttrString(enum_mod, "Enum");
sub_enum_type = PyObject_Call(enum_type, args, kwargs)
Py_DECREF(enum_type);
Py_DECREF(args);
Py_DECREF(kwargs);
return sub_enum_type