Page MenuHomec4science
Files F93261620

test_pickling.cpp
No OneTemporary
Actions

File Metadata

Created
Wed, Nov 27, 10:46

test_pickling.cpp
View Options

/*
tests/test_pickling.cpp -- pickle support
Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
All rights reserved. Use of this source code is governed by a
BSD-style license that can be found in the LICENSE file.
*/
#include "pybind11_tests.h"
TEST_SUBMODULE(pickling, m) {
// test_roundtrip
class Pickleable {
public:
Pickleable(const std::string &value) : m_value(value) { }
const std::string &value() const { return m_value; }
void setExtra1(int extra1) { m_extra1 = extra1; }
void setExtra2(int extra2) { m_extra2 = extra2; }
int extra1() const { return m_extra1; }
int extra2() const { return m_extra2; }
private:
std::string m_value;
int m_extra1 = 0;
int m_extra2 = 0;
};
class PickleableNew : public Pickleable {
public:
using Pickleable::Pickleable;
};
py::class_<Pickleable>(m, "Pickleable")
.def(py::init<std::string>())
.def("value", &Pickleable::value)
.def("extra1", &Pickleable::extra1)
.def("extra2", &Pickleable::extra2)
.def("setExtra1", &Pickleable::setExtra1)
.def("setExtra2", &Pickleable::setExtra2)
// For details on the methods below, refer to
// http://docs.python.org/3/library/pickle.html#pickling-class-instances
.def("__getstate__", [](const Pickleable &p) {
/* Return a tuple that fully encodes the state of the object */
return py::make_tuple(p.value(), p.extra1(), p.extra2());
})
.def("__setstate__", [](Pickleable &p, py::tuple t) {
if (t.size() != 3)
throw std::runtime_error("Invalid state!");
/* Invoke the constructor (need to use in-place version) */
new (&p) Pickleable(t[0].cast<std::string>());
/* Assign any additional state */
p.setExtra1(t[1].cast<int>());
p.setExtra2(t[2].cast<int>());
});
py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
.def(py::init<std::string>())
.def(py::pickle(
[](const PickleableNew &p) {
return py::make_tuple(p.value(), p.extra1(), p.extra2());
},
[](py::tuple t) {
if (t.size() != 3)
throw std::runtime_error("Invalid state!");
auto p = PickleableNew(t[0].cast<std::string>());
p.setExtra1(t[1].cast<int>());
p.setExtra2(t[2].cast<int>());
return p;
}
));
#if !defined(PYPY_VERSION)
// test_roundtrip_with_dict
class PickleableWithDict {
public:
PickleableWithDict(const std::string &value) : value(value) { }
std::string value;
int extra;
};
class PickleableWithDictNew : public PickleableWithDict {
public:
using PickleableWithDict::PickleableWithDict;
};
py::class_<PickleableWithDict>(m, "PickleableWithDict", py::dynamic_attr())
.def(py::init<std::string>())
.def_readwrite("value", &PickleableWithDict::value)
.def_readwrite("extra", &PickleableWithDict::extra)
.def("__getstate__", [](py::object self) {
/* Also include __dict__ in state */
return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
})
.def("__setstate__", [](py::object self, py::tuple t) {
if (t.size() != 3)
throw std::runtime_error("Invalid state!");
/* Cast and construct */
auto& p = self.cast<PickleableWithDict&>();
new (&p) PickleableWithDict(t[0].cast<std::string>());
/* Assign C++ state */
p.extra = t[1].cast<int>();
/* Assign Python state */
self.attr("__dict__") = t[2];
});
py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
.def(py::init<std::string>())
.def(py::pickle(
[](py::object self) {
return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
},
[](const py::tuple &t) {
if (t.size() != 3)
throw std::runtime_error("Invalid state!");
auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
cpp_state.extra = t[1].cast<int>();
auto py_state = t[2].cast<py::dict>();
return std::make_pair(cpp_state, py_state);
}
));
#endif
}

Event Timeline

c4science · Help