[QUESTION] Access self object using pickle factory (py::pickle).

I'm trying to implement a pickle serialization for a class hierarchy that includes Python-derived classes. I implemented this without using py::pickle (I defined directly the functions __getstate__ and __setstate__). However, I think this is deprecated and the new way to implement pickle support is through py::pickle. The problem with py::pickle is that it does not give access to py::object& self in the __setstate__ method.

This is my code:

// Abstract class
class A {
public:
    A(int d) : m_data(d) {}
    virtual ~A() {}
   
   // Serialize the common data for all the A classes.
    virtual py::tuple __getstate__() const {
        return py::make_tuple(m_data);
    }

    virtual int data() { return m_data; }

   // Some pure virtual method.
    virtual void pure_virtual() const = 0;
private:
// Some data that all the derived classes have.
    int m_data;
};

// Load the common data for A classes derived in C++.
template<typename DerivedA>
static std::shared_ptr<DerivedA> __setstate_Cpp__(py::tuple& t) {
    return std::make_shared<DerivedA>(t[0].cast<int>());
}

class CppA : public A {
public:
    using A::A;

    void pure_virtual() const override { }
};

// Trampoline for A.
class PyA : public A {
public:
    using A::A;

    virtual void pure_virtual() const override {
        PYBIND11_OVERRIDE_PURE(
            void,
            A,
            pure_virtual,
        );
    }

    virtual py::tuple __getstate__() const override {
       // get the common data for A.
        auto a_data = A::__getstate__();

        // get some extra data of the Python-derived class.
        pybind11::gil_scoped_acquire gil;
        pybind11::function override = pybind11::get_override(static_cast<const A*>(this), "__getstate_extra__");
        if (override) {
            auto o = override();
            // Includes some extra data.
            return py::make_tuple(a_data, true, py::make_tuple(o));
        }
       // Do not includes extra data.
        return py::make_tuple(a_data, false, py::make_tuple());
    }
    
  
    static py::object __setstate__(py::object& self, py::tuple& t) {
        // This __setstate__ function needs the Python-derived self-object to call the C++ initialization here!
        auto a_type = py::type::of<A>();
        auto a_data = t[0].cast<py::tuple>();
        // Initialize the C++ side!
        a_type.attr("__init__")(self, a_data[0].cast<int>());
        
        // Load the extra data for the Python-derived classes.
        bool has_extra_data = t[1].cast<bool>();
        if (has_extra_data) {
            auto extra_data = t[2].cast<py::tuple>();
            pybind11::gil_scoped_acquire gil;
            pybind11::function override = pybind11::get_override(self.cast<A*>(), "__setstate_extra__");
            if (override) {
                override(extra_data[0]);
            }
        }

        return self;
    }
};

PYBIND11_MODULE(example, m) {


    py::class_<A, PyA, std::shared_ptr<A>>(m, "A")
        .def(py::init<int>())
        .def("data", &A::data)
        .def("__getstate__", &A::__getstate__)
        .def("__setstate__", [](py::object& self, py::tuple& t){
            return PyA::__setstate__(self, t);
        });
        // .def(py::pickle(
        //     [](const A& a) {
        //         return a.__getstate__();
        //     },
        //     [](py::tuple& t) { <--- Here, I cannot access py::object& self!

        //     }
        // ));

    py::class_<CppA, A, std::shared_ptr<CppA>>(m, "CppA")
        .def(py::init<int>())
        .def(py::pickle(
            [](const CppA& self) {
                return self.__getstate__();
            },
            [](py::tuple& t) {  // In the C++ derived classes, I have to copy the setstate to avoid using the trampoline setstate
                return __setstate_Cpp__<CppA>(t);
            }
        ));
}

With this code, I can avoid implementing the serialization of common data for all the derived classes. It can be tested with the following code:

from example import *
import pickle

cpp_a = CppA(5)
bcpp = pickle.dumps(cpp_a)
loaded_cpp = pickle.loads(bcpp)

print("cpp_a data: " + str(cpp_a.data()))
print("loaded_cpp data: " + str(loaded_cpp.data()))

class PyDerived(A):
    def __init__(self, d):
        A.__init__(self, d)
        self.other_data = "other"

    def __getstate_extra__(self):
        return self.other_data

    def __setstate_extra__(self, d):
        self.other_data = d

py_a = PyDerived(6)
bpy = pickle.dumps(py_a)
loaded_py_a = pickle.loads(bpy)

print("py_a data: " + str(py_a.data()))
print("loaded_py_a data: " + str(loaded_py_a.data()))
print("extra_data: " + str(loaded_py_a.other_data))

this code correctly returns:

cpp_a data: 5
loaded_cpp data: 5
py_a data: 6
loaded_py_a data: 6
extra_data: other

The question is, how can I implement the same behavior using py::pickle? I do not have access to the py::object& self in __setstate__ using py::pickle.

The issue names no repository files or tests. Start by reading the py::pickle callback and the custom setstate lambda shown in the example, then verify the behavior with the provided Python pickle script; done means establishing whether py::pickle can provide the Python self object and documenting or implementing the supported approach.