dlib C++ Library - serialize_pickle.h

    // Copyright (C) 2013  Davis E. King (davis@dlib.net)
    // License: Boost Software License   See LICENSE.txt for the full license.
    #ifndef DLIB_SERIALIZE_PiCKLE_Hh_
    #define DLIB_SERIALIZE_PiCKLE_Hh_

    #include <dlib/serialize.h>
    #include <pybind11/pybind11.h>
    #include <sstream>
    #include <dlib/vectorstream.h>

    namespace py = pybind11;

    namespace dlib
    {

        template<typename T>
        py::tuple getstate(const T& item)
        {
            using namespace dlib;
            std::vector<char> buf;
            buf.reserve(5000);
            vectorstream sout(buf);
            serialize(item, sout);
            return py::make_tuple(py::handle(
                    PyBytes_FromStringAndSize(buf.size()?&buf[0]:0, buf.size())));
        }

        template<typename T>
        T setstate(py::tuple state)
        {
            using namespace dlib;
            if (len(state) != 1)
            {
                PyErr_SetObject(PyExc_ValueError,
                    py::str("expected 1-item tuple in call to __setstate__; got {}").format(state).ptr()
                );
                throw py::error_already_set();
            }

            // We used to serialize by converting to a str but the boost.python routines for
            // doing this don't work in Python 3.  You end up getting an error about invalid
            // UTF-8 encodings.  So instead we access the python C interface directly and use
            // bytes objects.  However, we keep the deserialization code that worked with str
            // for backwards compatibility with previously pickled files.
            T item;
            py::object obj = state[0];
            if (py::isinstance<py::str>(obj))
            {
                py::str data = state[0].cast<py::str>();
                std::string temp = data;
                std::istringstream sin(temp);
                deserialize(item, sin);
            }
            else if(PyBytes_Check(py::object(state[0]).ptr()))
            {
                py::object obj = state[0];
                char* data = PyBytes_AsString(obj.ptr());
                unsigned long num = PyBytes_Size(obj.ptr());
                std::istringstream sin(std::string(data, num));
                deserialize(item, sin);
            }
            else
            {
                throw error("Unable to unpickle, error in input file.");
            }

            return item;
        }

    }

    #endif // DLIB_SERIALIZE_PiCKLE_Hh_