在 C 扩展中定义 Python 枚举 - 我这样做对吗？

Question

我正在开发一个 Python C 扩展，我想公开一个完全在 C 中定义的自定义枚举（如：继承自 enum.Enum 的类）。

事实证明这不是一个微不足道的任务，并且继承使用的常规机制.tp_base不起作用 - 很可能是由于 Enum 的元类没有被引入。

基本上我正在尝试这样做：

import enum

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

在C.

经过大量深入研究 cpython 的内部结构后，这就是我想到的，包装在一个示例可构建模块中：

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

基本上，我直接调用EnumMetas__prepare__方法来创建正确的 classdict，然后我PyType_Type也调用该对象来创建子类型。

这是可行的，并且 AFAICT 会产生一个行为完全符合预期的类，但是......我这样做对吗？如有任何反馈，我们将不胜感激。

Answer 1

metaclass是的，这Enum很棘手。

但你可以在这里看到你可以创建一个枚举（在 Python 中），如下所示：

FooBar = enum.Enum('FooBar', dict(FOO=1, BAR=2))

因此，您可以使用此技术在 Python C-API 中轻松创建枚举类，方法如下：

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