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

[Bartosz Golaszewski]

Hi!

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.

I want to use this in a real project, namely the v2 of libgpiod python
bindings[1] so it's important to me to get this right.

Best regards,
Bartosz Golaszewski

[1] https://git.kernel.org/pub/scm/libs/libgpiod/libgpiod.git/

[MRAB]

On 2021-07-23 09:20, Bartosz Golaszewski wrote:
> Hi!
>
> 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:
>

[snip]

>
> 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. */

PyTuple_Pack doesn't steal references, as far as I can tell.

> if (!bases)
> Py_DECREF(enum_type);
>
> return bases;
> }
>

[snip]

[Bartosz Golaszewski]

Right, the doc says it's equivalent to Py_BuildValue("(OO...)", ...)
and it does increase the reference count on stored objects. It doesn't
answer the main question though. :)

Bartosz

> > if (!bases)
> > Py_DECREF(enum_type);
> >
> > return bases;
> > }
> >
> [snip]

> --
> https://mail.python.org/mailman/listinfo/python-list

[Dan Stromberg]

On Fri, Jul 23, 2021 at 1:20 AM Bartosz Golaszewski <br...@bgdev.pl> wrote:

> Hi!
>
> 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'm probably missing something obvious, but why would you write new code in
C when you can just use Cython? Cython is a lot easier, and quite fast,
and should (eventually?) allow compiling to HPY instead of just "the" C
extension module interface.

https://news.ycombinator.com/item?id=26627683

[Bartosz Golaszewski]

I'm the author and maintainer of libgpiod - the user-space library and
tools for using the linux GPIO character device.

The core library is written in C but we're also exposing C++ and
Python bindings (with more language bindings planned). The python
bindings are written as a C extension module and look like this:
https://git.kernel.org/pub/scm/libs/libgpiod/libgpiod.git/tree/bindings/python/gpiodmodule.c.

We're in the process of writing the (backward incompatible) version 2
of the library in order to support the new kernel features and the C
API has changed a lot so we're also rewriting the bindings. Among
others: all bitwise flags have now been converted to enums, hence my
question. In C++ we'll use scoped enum classes and I'd like to do a
similar thing in python.

What I'm doing is not aimed at using C for speed but for calling the C
APIs. I know I could use SWIG but in order to make the interface
elegant, it would have to be packaged in proper Python classes anyway,
creating another layer of code so I prefer to just use C.

Bart

[Serhiy Storchaka]

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.

[Schachner, Joseph]

Instead of struggling to define an enum in C that can be read in Python - I'm assuming you can pass strings back and forth - why not just print whatever you need to give to Python into a string (or maybe 2 strings) and send it to Python as string? Python is a dynamic language, it can quickly build an enum type from one string and instantiate enums of that type from the other string.

In short, if you have problems creating an enum in C and passing it to Python, give the problems to Python! Let it create the enum.

--- Joseph S.


Teledyne Confidential; Commercially Sensitive Business Data

[Bartosz Golaszewski]

On Fri, Jul 30, 2021 at 2:41 PM Serhiy Storchaka <stor...@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

[Bartosz Golaszewski]

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

[Sean DiZazzo]

No

[Serhiy Storchaka]

03.08.21 13:03, Bartosz Golaszewski пише:

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

I think it can be simpler.

1. No need to create the __main__ module. You can just create a dict. If
some attributes are required (e.g. __name__) it is easy to set them in
the Python code (__name__ = 'pycenum').

2. No need to import the enum module in the C code. It is easier to do
it in the Python code.

3. Can you define your DummyType in the Python code instead of the C
code? Even if you need to create in the C code, it may be better to
define your enum class inside a dummy DummyType. It will set correct
__qualname__ of the enum class:

__name__ = 'pycenum'
import enum
class DummyType:
class FooBar(enum.Enum):
...

4. Did you consider idea of making DummyType a heap-allocated type? It
may be easy to create new type with PyType_FromModuleAndSpec(), and it
is more stable API, and the type is not immutable, so it is easy to add
new attributes with PyObject_SetAttrString().

[Bartosz Golaszewski]

On Fri, Aug 6, 2021 at 9:02 PM Serhiy Storchaka <stor...@gmail.com> wrote:
>
> 03.08.21 13:03, Bartosz Golaszewski пише:
> > Just a follow-up: this is how I did it eventually:
>
> I think it can be simpler.
>
> 1. No need to create the __main__ module. You can just create a dict. If
> some attributes are required (e.g. __name__) it is easy to set them in
> the Python code (__name__ = 'pycenum').
>
> 2. No need to import the enum module in the C code. It is easier to do
> it in the Python code.
>

Don't you need the '__import__' attribute for that? It's linked to point #1.

> 3. Can you define your DummyType in the Python code instead of the C
> code? Even if you need to create in the C code, it may be better to
> define your enum class inside a dummy DummyType. It will set correct
> __qualname__ of the enum class:
>

This code is a PoC for the real code I want to use it in where the
type embedding the enum is defined in C.

> __name__ = 'pycenum'
> import enum
> class DummyType:
> class FooBar(enum.Enum):
> ...
>
> 4. Did you consider idea of making DummyType a heap-allocated type? It
> may be easy to create new type with PyType_FromModuleAndSpec(), and it
> is more stable API, and the type is not immutable, so it is easy to add
> new attributes with PyObject_SetAttrString().
>

Makes sense, thanks, I'll try it.

Bart