How to create an Enum object in Python C API?

Question

I'm struggling how to create a python Enum object inside the Python C API. The enum class has assigned tp_base to PyEnum_Type, so it inherits Enum. But, I can't figure out a way to tell the Enum base class what items are in the enum. I want to allow iteration and lookup from Python using the __members__ attribute that every Python Enum provides.

Thank you,

Jelle

Answer 1

It is not straightforward at all. The Enum is a Python class using a Python metaclass. It is possible to create it in C but it will be just emulating the constructing Python code in C - the end result is the same and while it speeds up things slightly, you'll most probably run the code only once within each program run.

In any case it is possible, but it is not easy at all. I'll show how to do it in Python:

from enum import Enum

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

print(Color)
print(Color.RED)

is the same as:

from enum import Enum

name = 'Color'
bases = (Enum,)
enum_meta = type(Enum)

namespace = enum_meta.__prepare__(name, bases)
namespace['RED'] = 1
namespace['GREEN'] = 2
namespace['BLUE'] = 3

Color = enum_meta(name, bases, namespace)

print(Color)
print(Color.RED)

The latter is the code that you need to translate into C.

Answer 2

Edited note: An answer on a very similar question details how enum.Enum has a functional interface that can be used instead. That is almost certainly the correct approach. I think my answer here is a useful alternative approach to be aware of, although it probably isn't the best solution to this problem.

---

I'm aware that this answer is slightly cheating, but this is exactly the kind of code that's better written in Python, and in the C API we still have access to the full Python interpreter. My reasoning for this is that the main reason to keep things entirely in C is performance, and it seems unlikely that creating enum objects will be performance critical.

I'll give three versions, essentially depending on the level of complexity.

---

First, the simplest case: the enum is entirely known and defined and compile-time. Here we simply set up an empty global dict, run the Python code, then extract the enum from the global dict:

PyObject* get_enum(void) {
    const char str[] = "from enum import Enum\n"
                       "class Colour(Enum):\n"
                       "    RED = 1\n"
                       "    GREEN = 2\n"
                       "    BLUE = 3\n"
                       "";
    PyObject *global_dict=NULL, *should_be_none=NULL, *output=NULL;
    global_dict = PyDict_New();
    if (!global_dict) goto cleanup;
    should_be_none = PyRun_String(str, Py_file_input, global_dict, global_dict);
    if (!should_be_none) goto cleanup;
    // extract Color from global_dict
    output = PyDict_GetItemString(global_dict, "Colour");
    if (!output) {
        // PyDict_GetItemString does not set exceptions
        PyErr_SetString(PyExc_KeyError, "could not get 'Colour'");
    } else {
        Py_INCREF(output); // PyDict_GetItemString returns a borrow reference
    }
    cleanup:
    Py_XDECREF(global_dict);
    Py_XDECREF(should_be_none);
    return output;
}

---

Second, we might want to change what we define in C at runtime. For example, maybe the input parameters pick the enum values. Here, I'm going to use string formatting to insert the appropriate values into our string. There's a number of options here: sprintf, PyBytes_Format, the C++ standard library, using Python strings (perhaps with another call into Python code?). Pick whichever you're most comfortable with.

PyObject* get_enum_fmt(int red, int green, int blue) {
    const char str[] = "from enum import Enum\n"
                       "class Colour(Enum):\n"
                       "    RED = %d\n"
                       "    GREEN = %d\n"
                       "    BLUE = %d\n"
                       "";
    PyObject *formatted_str=NULL, *global_dict=NULL, *should_be_none=NULL, *output=NULL;

    formatted_str = PyBytes_FromFormat(str, red, green, blue);
    if (!formatted_str) goto cleanup;
    global_dict = PyDict_New();
    if (!global_dict) goto cleanup;
    should_be_none = PyRun_String(PyBytes_AsString(formatted_str), Py_file_input, global_dict, global_dict);
    if (!should_be_none) goto cleanup;
    // extract Color from global_dict
    output = PyDict_GetItemString(global_dict, "Colour");
    if (!output) {
        // PyDict_GetItemString does not set exceptions
        PyErr_SetString(PyExc_KeyError, "could not get 'Colour'");
    } else {
        Py_INCREF(output); // PyDict_GetItemString returns a borrow reference
    }
    cleanup:
    Py_XDECREF(formatted_str);
    Py_XDECREF(global_dict);
    Py_XDECREF(should_be_none);
    return output;
}

Obviously you can do as much or as little as you like with string formatting - I've just picked a simple example to show the point. The main differences from the previous version are the call to PyBytes_FromFormat to set up the string, and the call to PyBytes_AsString that gets the underlying char* out of the prepared bytes object.

---

Finally, we could prepare the enum attributes in C Python dict and pass it in. This necessitates a bit of a change. Essentially I use @AnttiHaapala's lower-level Python code, but insert namespace.update(contents) after the call to __prepare__.

PyObject* get_enum_dict(const char* key1, int value1, const char* key2, int value2) {
    const char str[] = "from enum import Enum\n"
                       "name = 'Colour'\n"
                       "bases = (Enum,)\n"
                       "enum_meta = type(Enum)\n"
                       "namespace = enum_meta.__prepare__(name, bases)\n"
                       "namespace.update(contents)\n"
                       "Colour = enum_meta(name, bases, namespace)\n";

    PyObject *global_dict=NULL, *contents_dict=NULL, *value_as_object=NULL, *should_be_none=NULL, *output=NULL;
    global_dict = PyDict_New();
    if (!global_dict) goto cleanup;

    // create and fill the contents dictionary
    contents_dict = PyDict_New();
    if (!contents_dict) goto cleanup;
    value_as_object = PyLong_FromLong(value1);
    if (!value_as_object) goto cleanup;
    int set_item_result = PyDict_SetItemString(contents_dict, key1, value_as_object);
    Py_CLEAR(value_as_object);
    if (set_item_result!=0) goto cleanup;
    value_as_object = PyLong_FromLong(value2);
    if (!value_as_object) goto cleanup;
    set_item_result = PyDict_SetItemString(contents_dict, key2, value_as_object);
    Py_CLEAR(value_as_object);
    if (set_item_result!=0) goto cleanup;

    set_item_result = PyDict_SetItemString(global_dict, "contents", contents_dict);
    if (set_item_result!=0) goto cleanup;

    should_be_none = PyRun_String(str, Py_file_input, global_dict, global_dict);
    if (!should_be_none) goto cleanup;
    // extract Color from global_dict
    output = PyDict_GetItemString(global_dict, "Colour");
    if (!output) {
        // PyDict_GetItemString does not set exceptions
        PyErr_SetString(PyExc_KeyError, "could not get 'Colour'");
    } else {
        Py_INCREF(output); // PyDict_GetItemString returns a borrow reference
    }
    cleanup:
    Py_XDECREF(contents_dict);
    Py_XDECREF(global_dict);
    Py_XDECREF(should_be_none);
    return output;
}

Again, this presents a reasonably flexible way to get values from C into a generated enum.

---

For the sake of testing I used the follow simple Cython wrapper - this is just presented for completeness to help people try these functions.

cdef extern from "cenum.c":
    object get_enum()
    object get_enum_fmt(int, int, int)
    object get_enum_dict(char*, int, char*, int)


def py_get_enum():
    return get_enum()

def py_get_enum_fmt(red, green, blue):
    return get_enum_fmt(red, green, blue)

def py_get_enum_dict(key1, value1, key2, value2):
    return get_enum_dict(key1, value1, key2, value2)

---

To reiterate: this answer is only partly in the C API, but the approach of calling Python from C is one that I've found productive at times for "run-once" code that would be tricky to write entirely in C.