Search code examples
c++enumspybind11

How to support len() method for Python enums created by Pybind11


Suppose that I have a C++ enumeration like this:

enum class Kind {Kind1 = 1, Kind2, Kind3};

to bind this enumeration into a Python enumeration using Pybind11, I am doing something like this:

py::enum_<Kind>(py_module, "Kind")
   .value("Kind1", Kind::Kind1)
   .value("Kind2", Kind::Kind2)
   .value("Kind3", Kind::Kind3)
   .def("__len__", 
      [](Kind p) {
         return 3;
      });

After compiling the code, if I ask for the length of the enumration, I will get this error:

>>> len(Kind)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: object of type 'pybind11_type' has no len()

Any ideas how I can fix it?

Edit 1: I am using Pybind11 version 2.10.1 on Visual Studio 2019 (C++17).

Edit 2: I want to have the same behavior as it is in Python Enums:

>>> from enum import Enum
>>> class Kind(Enum):
...    kind1 = 1
...    kind2 = 2
...    kind3 = 3
...
>>> len(Kind)
3

Solution

  • Before we begin, let's solve one related problem -- how can we avoid having to explicitly specify the length to return for each enumeration? It would be nice if something did it for us, or perhaps if we could calculate it dynamically during runtime.

    Turns out there is a way. While reading through the code, I came across an interesting implementation detail. The enumeration wrapper class that pybind11 creates has an attribute named __entries. It is a dictionary, which holds one entry per enumeration value, mostly used to generate documentation, get text representation of the value, and export the values to parent scope.

    Here is what it looks like for your example enum:

    >>> print(Kind.__entries)
    {'Kind1': (Kind.Kind1, None), 'Kind2': (Kind.Kind2, None), 'Kind3': (Kind.Kind3, None)}
    

    Hence, we can use len(Kind.__entries) to get the correct length (number of enumeration values) at runtime. In C++ this would be py::len(cls.attr("__entries")) where cls is the Kind class object.


    Now we can get to the root of the issue -- how can we make len work on a class object, rather than a class instance. According to this SO answer, one way to accomplish that is using a metaclass. Specifically, we need the enum wrapper class to use a metaclass that has a __len__ member function, which will calculate and return the number of values the wrapper class holds.

    It turns out that the wrapper classes generated by pybind already use a custom metaclass named pybind11_type:

    >>> type(Kind)
    <class 'pybind11_builtins.pybind11_type'>
    

    Hence, the approach would be to create a new metaclass, say pybind11_ext_enum, which would be derived from pybind11_type, and provide the missing __len__.


    The next question is, how can we create such metaclass from c++. Pybind11 doesn't provide any convenience functionality to do this, so we'll have to do it ourselves. To do so, we need:

    1. An object representing the original pybind11 metaclass pybind11_type. I found it stashed in internals, so I grab it from there.

      py::object get_pybind11_metaclass()
      {
          auto &internals = py::detail::get_internals();
          return py::reinterpret_borrow<py::object>((PyObject*)internals.default_metaclass);
      }
      
    2. An object representing the standard Python metaclass type (i.e. PyType_Type in CPython API).

      py::object get_standard_metaclass()
      {
          auto &internals = py::detail::get_internals();
          return py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
      }
      
    3. A dictionary of attributes we want this new class to have. This needs only a single entry, defining our __len__ method.

      py::dict attributes;
      attributes["__len__"] = py::cpp_function(
          [](py::object cls) {
              return py::len(cls.attr("__entries"));
          }
          , py::is_method(py::none())
          );
      
    4. Use type to create our new class object.

      auto pybind11_metaclass = get_pybind11_metaclass();
      auto standard_metaclass = get_standard_metaclass();
      return standard_metaclass(std::string("pybind11_ext_enum")
          , py::make_tuple(pybind11_metaclass)
          , attributes);
      

    We can put parts 3 and 4 into a function: py::object create_enum_metaclass() { ... }.


    Finally, we have to use our new metaclass when creating the enum wrapper.

    PYBIND11_MODULE(so07, m)
    {
        auto enum_metaclass = create_enum_metaclass();
        py::enum_<Kind>(m, "Kind", py::metaclass(enum_metaclass))
            .value("Kind1", Kind::Kind1)
            .value("Kind2", Kind::Kind2)
            .value("Kind3", Kind::Kind3)
            ;
    }
    

    And now we can use it in Python:

    >>> from so07 import Kind
    >>> type(Kind)
    <class 'importlib._bootstrap.pybind11_ext_enum'>
    >>> len(Kind)
    3