Defining a Python enum in a C extension - am I doing this right?

[Bartosz Golaszewski]

On Sat, Jul 31, 2021 at 3:01 PM Bartosz Golaszewski <brgl at bgdev.pl> wrote:
>
> On Fri, Jul 30, 2021 at 2:41 PM Serhiy Storchaka <storchaka at gmail.com> wrote:
> >
> > 23.07.21 11:20, Bartosz Golaszewski пише:
> > > 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.
> >
> > I think that it would be much easier to define it in Python, and then
> > either import a Python module in your C code, or exec a Python code as a
> > string.
> >
>
> You mean: evaluate a string like this:
>
> '''
> import enum
>
> class FooBar(enum.Enum):
>     FOO = 1
>     BAR = 2
>     BAZ = 3
> '''
>
> And then pull in the FooBar type from the resulting dictionary into my
> C code? Sounds good actually. I think I'll be able to add the FooBar
> type to another type's tp_dict too - because some enums I want to
> create will be nested in other classes.
>
> Bart

Just a follow-up: this is how I did it eventually:

```
#include <Python.h>

typedef struct {
    PyObject_HEAD;
} dummy_object;

PyDoc_STRVAR(dummy_type_doc, "Dummy type in which the enum will be nested.");

static PyTypeObject dummy_type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "pycenum.DummyType",
    .tp_basicsize = sizeof(dummy_object),
    .tp_flags = Py_TPFLAGS_DEFAULT,
    .tp_doc = dummy_type_doc,
    .tp_new = PyType_GenericNew,
    .tp_dealloc = (destructor)PyObject_Del,
};

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

static int add_foobar_enum(PyObject *module)
{
    static const char *foobar_src =
        "class FooBar(enum.Enum):\n"
        "    FOO = 1\n"
        "    BAR = 2\n"
        "    BAZ = 3\n";

    PyObject *main_mod, *main_dict, *enum_mod, *result, *foobar_type;
    int ret;

    main_mod = PyImport_AddModule("__main__");
    if (!main_mod)
        return -1;

    main_dict = PyModule_GetDict(main_mod);
    if (!main_dict) {
        Py_DECREF(main_mod);
        return -1;
    }

    enum_mod = PyImport_ImportModule("enum");
    if (!enum_mod) {
        Py_DECREF(main_mod);
        return -1;
    }

    ret = PyDict_SetItemString(main_dict, "enum", enum_mod);
    Py_DECREF(enum_mod);
    if (ret) {
        Py_DECREF(main_mod);
        return -1;
    }

    result = PyRun_String(foobar_src, Py_single_input,
                  main_dict, main_dict);
    if (!result) {
        Py_DECREF(main_mod);
        return -1;
    }

    foobar_type = PyDict_GetItemString(main_dict, "FooBar");
    if (!foobar_type) {
        Py_DECREF(main_mod);
        return -1;
    }

    ret = PyDict_SetItemString(dummy_type.tp_dict, "FooBar", foobar_type);
    Py_DECREF(foobar_type);
    Py_DECREF(main_mod);
    if (ret)
        return -1;

    PyType_Modified(&dummy_type);

    return ret;
}

PyMODINIT_FUNC PyInit_pycenum(void)
{
    PyObject *module;
    int ret;

    module = PyModule_Create(&module_def);
    if (!module)
        return NULL;

    ret = PyModule_AddStringConstant(module, "__version__", "0.0.1");
    if (ret) {
        Py_DECREF(module);
        return NULL;
    }

    ret = PyType_Ready(&dummy_type);
    if (ret) {
        Py_DECREF(module);
        return NULL;
    }

    ret = add_foobar_enum(module);
    if (ret) {
        Py_DECREF(module);
        return NULL;
    }

    Py_INCREF(&dummy_type);
    ret = PyModule_AddObject(module, "DummyType", (PyObject *)&dummy_type);
    if (ret) {
        Py_DECREF(&dummy_type);
        Py_DECREF(module);
        return NULL;
    }

    return module;
}
```

Bart