Dominick Pastore
Dominick Pastore

Reputation: 4435

How to return a Python Enum from C extension?

I'm writing a C extension for Python that wraps a C library. The C library has a few enum types and I've written corresponding IntEnums for these, for example:

from enum import IntEnum
# _enum_consts is a C extension creating Python
# constants for the library's enum constants
from . import _enum_consts

class Align(IntEnum):
    DEFAULT = _enum_consts.MD_ALIGN_DEFAULT
    LEFT = _enum_consts.MD_ALIGN_LEFT
    CENTER = _enum_consts.MD_ALIGN_CENTER
    RIGHT = _enum_consts.MD_ALIGN_RIGHT

Now, in the C extension where I'm wrapping the library's main functionality, some of the library functions return enum types, and I want to convert those to the Python IntEnum values. (I realize I could simply return the ints directly and they would successfully compare to IntEnums, but that sort of defeats the purpose.)

The question is, how do I create the IntEnum values to return? All I can find are examples where the C enum constants are converted to Python constants with PyModule_AddIntConstant() (which is just what the _enum_consts extension module does).

Upvotes: 4

Views: 1422

Answers (1)

Dominick Pastore
Dominick Pastore

Reputation: 4435

First things first, the module with the IntEnum needs to be imported into the extension module (if you are not already doing that). In the module initialization function:

PyMODINIT_FUNC PyInit_someextension(void)
{
    // ...

    PyObject *m;
    m = PyModule_Create(&someextension_module);
    if (m == NULL) {
        return NULL;
    }

    // ...

    // Import the enums module, where "somepackage.enums"
    // is the full name of the enums module
    PyObject *enums = PyImport_ImportModule("somepackage.enums");
    if (enums == NULL) {
        Py_DECREF(m);
        return NULL;
    }
    Py_DECREF(enums);

    return m;
}

Then, where you want to instantiate an IntEnum value, first fetch the module it's in again:

PyObject *enums = PyImport_AddModule("somepackage.enums");
if (enums == NULL) {
    // Return NULL, -1, etc. to signal an exception
}

Right after, fetch your IntEnum type (Align in this case) from the module:

PyObject *align_enum = PyObject_GetAttrString(enums, "Align");
if (align_enum == NULL) {
    // Return NULL, -1, etc. to signal an exception
}

Then, you can create the instance you need. align here will be the value to fetch from the IntEnum. We do this by calling the type object, just as we would in regular Python (note that PyObject_CallFunction is not actually limited to calling functions--it can call any callable object):

PyObject *align_instance = PyObject_CallFunction(align_enum, "(i)", align);
if (align_instance == NULL) {
    Py_DECREF(align_enum);
    // Return NULL, -1, etc. to signal an exception
}

Now, align_instance contains the IntEnum instance to return. Be sure to Py_DECREF(align_enum) before returning.

Upvotes: 2

Related Questions