How to make my SWIG extension module work with Pickle?

Question

I have an extension module for Python that uses SWIG as a wrapper and I try to serialize it with Pickle and I fail =)

1. If anyone has a source of SWIG extension

Answer 1

Here are a few additional methods. None is of as general applicability as the accepted answer is, but if your class meets some (simple) requirements then you can make pickling easier on your users by making the instances themselves (not wrapped versions) picklable. These techniques are all used by the LSST afw package.

Note that when unpickling using the __getstate__/__setstate__ pair, the __init__ method will not be called, which means that unless you're careful, you'll have an object that you can't do anything with (if you keep getting NotImplementedError: Wrong number or type of arguments for overloaded function, this is a possibility). This drives us to use __reduce__ (or you could call __init__ from __setstate__).

If you're SWIG-ing class Foo that takes constructor arguments that you have access to from the instance (e.g., via accessors), add the following to your interface (.i) file:

%extend Foo {
%pythoncode {
    def __reduce__(self):
        # Requires matching constructor: __init__(foo, bar)
        args = self.getFoo(), self.getBar()
        return self.__class__, args
}
}

If you can create your object with a default constructor and then manipulate it to regain its former state, use something like this:

%extend Foo {
%pythoncode {
    def __getstate__(self):
        args = self.getFoo(), self.getBar()
        return args
    def __setstate__(self, state):
        # Requires empty constructor: __init__()
        self.__init__()
        foo, bar = state
        self.setFoo(foo)
        self.setBar(bar)
}
}

Alternatively, if your class can do a serialisation of binary data to/from memory (e.g., some in-memory representation of your own on-disk format):

%include "cdata.i"

%extend Foo {
%pythoncode {
    def __reduce__(self):
        s = Serializer()
        self.serialize(s)
        size = s.getLength()
        data = cdata(s.getData(), size)
        return unreduceFoo, (data, size)
}
}

%pythoncode {
def unreduceFoo(data, size):
    s = Serializer(size)
    memmove(s.getData(), data)
    return Foo(s)
}

Finally, if you're using boost::serialization, use this snippet by Sogo Mineo:

%{
    #include <boost/serialization/serialization.hpp>
    #include <boost/archive/binary_oarchive.hpp>
    #include <boost/archive/binary_iarchive.hpp>
    #include <sstream>
%}
%include "std_string.i"

%define %boost_picklable(cls...)
    %extend cls {
        std::string __getstate__()
        {
            std::stringstream ss;
            boost::archive::binary_oarchive ar(ss);
            ar << *($self);
            return ss.str();
        }

        void __setstate_internal(std::string const& sState)
        {
            std::stringstream ss(sState);
            boost::archive::binary_iarchive ar(ss);
            ar >> *($self);
        }

        %pythoncode %{
            def __setstate__(self, sState):
                self.__init__()
                self.__setstate_internal(sState)
        %}
    }
%enddef

%boost_picklable(Foo)

Answer 2

I had to make a minor change to the accepted answer to make it work for my case. Since the initialization function of my class has some input arguments, I had to add an extra argument *arg to the C.__init(self) function, like so:

    class PickalableC(C, PickalableSWIG):
        def __init__(self, *args):
            self.args = args
            C.__init__(self, *args)

Maybe this is helpful for someone. I hope it is not too trivial.

I posted it as an answer since I cannot comment and learned it is not something you edit a post for.

Answer 3

Seems like I found simlple solution that works for me:

So let's say we have class C that was generated with SWIG, then we wrap it with

class PickalableC(C, PickalableSWIG):

    def __init__(self, *args):
        self.args = args
        C.__init__(self)

where PickalableSWIG is

class PickalableSWIG:

    def __setstate__(self, state):
        self.__init__(*state['args'])

    def __getstate__(self):
        return {'args': self.args}

Then

pickle.loads(pickle.dumps(C()))

fails, but

pickle.loads(pickle.dumps(PickalableC()))

succeeds =)