Python Protocol-like support with implicitly_convertible

[bpdavis86]

I wish something like this would have been in the documentation because it took me some time to find just the right way to go about it.

We are attempting a mixed Python/C++ simulation project with the goal of interchangeable components from both sides. There are 4 use cases.

1. Pure Python Simulation
2. Pure C++ Simulation
3. Python main sim with C++ plug-in entities
4. C++ main sim with Python plug-in entities (but possibly invoked from Python interpreter)

Additionally, both the pure Python and pure C++ simulation should be runnable without the involvement of pybind, so the implementation of the wrappers should be decoupled.

The first two cases are obviously straightforward. The third is not too hard, because we can pybind the C++ modules, and as long as they fulfil the correct interface, we do not care. Python has a nice feature "protocols" to motivate this.

from typing import Protocol

class ISimulation(Protocol):
    def add_entity(self, entity: IEntity) -> None: ...

    def run() -> None: ...

class IEntity(Protocol):
    def update_time(self, time: float) -> None: ...

    @property
    def name(self) -> str: ...

Then, if I implement ISimulation in Python following these interfaces, I do not care what kind of entity comes in, it can be a Python entity, or C++ pybind entity that has the right functions. I can swap these things in and out in loader code.

To do things the other way round, it is harder. I can make an IEntity virtual base class in C++ easily enough.

class ISimulation
{
    virtual void add_entity(std::shared_ptr<IEntity> entity) = 0;
    virtual void run() = 0;
};
class IEntity
{
    virtual void update_time(double time) = 0;
    virtual std::string name() const = 0;
};

PYBIND11_MODULE(cppsim, m)
{
py::class_<ISimulation, std::shared_ptr<ISimulation>>(m, "ISimulation");
    .def("add_entity", &ISimulation::add_entity)
    .def("run", &ISimulation::run);

py::class_<IEntity, std::shared_ptr<IEntity>>(m, "IEntity")
    .def("update_time", &IEntity::update_time)
    .def_property_readonly("name", &IEntity::name);
}

Then to make this work in Python by the standard documentation, one needs to derive Python implementations from the pybind IEntity class, and add associated "trampoline" class. But, I don't want to do this because it wrecks the concept of Protocol in Python and the Python-only use case 1 (which should not have to load anything pybind-related).

So, I would like to achieve something like "Protocol" in pybind, i.e., I can accept an arbitrary Python object which implicitly fulfils a given interface and have it made compatible with my strongly typed C++ simulation object.

Obviously we will need a "real" C++ derived class of IEntity to make this work. After much meandering, I came up with this:

// Helper metafunction to tell us what kind of wrapper implementation class to use for a given interface type
// wrapper_impl<Interface>::type will tell us the class
template <typename Interface>
struct wrapper_impl
{
};

template <typename Interface>
class PythonWrapper : public Interface
{
public:
    // get the underlying Python object
    py::object obj() { return _obj; }

    // If the Python object already has a C++ wrapper attached, get that.
    // Otherwise, attempt to make a wrapper of this type over top of the Python object.
    // The wrapper will then be attached to the Python object to keep it alive.
    // Note that this requires the input Python object support dynamic properties.
    static std::shared_ptr<Interface> get(py::object obj)
    {
        std::shared_ptr<Interface> ptr;

        try
        {
            // try to get existing wrapper object
            ptr = obj.attr("_cpp_obj").cast<std::shared_ptr<Interface>>();
        }
        catch (const py::error_already_set &e)
        {
            if (!e.matches(PyExc_AttributeError))
            {
                throw;
            }
        }
        catch (const py::cast_error &e)
        {
        }

        if (!ptr)
        {
            // Try to make a new wrapper and attach to the Python object.
            // If the interface does not match, we'll get attribute errors here.
            ptr = std::make_shared<wrapper_impl<Interface>::type>(obj);
            obj.attr("_cpp_obj") = ptr;
        }

        return ptr;
    }

    // Since we hold a Python object internally, we need to support GC slots.
    // This code to be attached to the class binding.
    static void type_setup(PyHeapTypeObject *heap_type)
    {
        auto *type = &heap_type->ht_type;
        type->tp_flags |= Py_TPFLAGS_HAVE_GC;
        type->tp_traverse = [](PyObject *self_base, visitproc visit, void *arg)
        {
            auto &self = py::cast<wrapper_impl<Interface>::type &>(py::handle(self_base));
            Py_VISIT(self.obj().ptr());
            return 0;
        };
        type->tp_clear = [](PyObject *self_base)
        {
            auto &self = py::cast<wrapper_impl<Interface>::type &>(py::handle(self_base));
            self._obj = py::none();
            return 0;
        };
    }

protected:
    PythonWrapper(py::object obj) : _obj(obj) {}
    py::object _obj;
};

// Instantiate a wrapper for our entity interface
class PyEntity : public PythonWrapper<IEntity>
{
public:
    PyEntity(py::object obj) : PythonWrapper(obj) { _initialize_functions(); }

    // virtual functions simply forward to captures of the Python attributes
    virtual std::string name() const override { return _name(); }
    virtual void update_time(double time) override { return _update_time(time); }

private:
    std::function<std::string()> _name;
    std::function<void(double)> _update_time;

    void _initialize_functions()
    {
        // Go ahead and load all the Python attributes into the wrapper functions at construction.
        // This way, if the interface is broken, we "fail early".
        py::type py_type = py::type::of(_obj);

        py::function getter;
        // for a property, we have to get the function from the type object
        getter = py_type.attr("name").attr("fget");
        _name = [getter, this]() -> std::string
        {
            return getter(_obj).cast<std::string>();
        };

        py::function update_time_py = _obj.attr("update_time");
        _update_time = [update_time_py](double time)
        { return update_time_py(time); };
    }
};

// Register our implementation with the metafunction
template <>
struct wrapper_impl<IEntity>
{
    using type = PyEntity;
};

Now, this can be bound as follows.

py::class_<IEntity, std::shared_ptr<IEntity>>(m, "IEntity")
    // Provide Python-side constructor which creates the wrapper type
    .def(py::init([](py::object obj) { return PythonWrapper<IEntity>get(obj); )
    .def("update_time", &IEntity::update_time)
    .def_property_readonly("name", &IEntity::name);

// Key, we can convert directly from objects to IEntity if the interface works
py::implicitly_convertible<py::object, IEntity>();

// Need to bind the wrapper class so the GC setup is properly taken care of.
// Also, provide the interface to get the Python object back out.
py::class_<PyEntity, IEntity, std::shared_ptr<PyEntity>>(m, "PyEntity", py::custom_type_setup(&PyEntity::type_setup))
    .def(py::init<py::object>())
    .def_property_readonly("obj", &PyEntity::obj);

// assume you wrote a concrete Simulation in C++
py::class_<Simulation, ISimulation, std::shared_ptr<Simulation>>(m, "Simulation")
    // bind your concrete C++ simulation
    ;

We can then do something like this:

from cppsim import Simulation

class MyPyEntity:
    def update_time(time):
        # do things
        pass

    @property
    def name(self):
        return "My Entity"

sim = Simulation()
e = MyPyEntity()
sim.add_entity(e)
sim.run()

This will work, even though MyPyEntity only follows the IEntity "Protocol" rather than being derived from the bound IEntity class. This happens without any explicit wrapping in a C++ type thanks to implicitly_convertible. The only restriction is that the input Python object must support dynamic attributes (so it can be made to own its C++ wrapper). Perhaps there is some clever way around this.

It is a little bit laborious to write the Python attribute dispatch wrappers if you have a large interface, but the result is very nice. You can also cache things in the wrapper object if you have e.g. an API with reference types.