Wenzel Jakob | 1c329aa | 2016-04-13 02:37:36 +0200 | [diff] [blame] | 1 | /* |
Dean Moldovan | a0c1ccf | 2016-08-12 13:50:00 +0200 | [diff] [blame] | 2 | tests/test_pickling.cpp -- pickle support |
Wenzel Jakob | 1c329aa | 2016-04-13 02:37:36 +0200 | [diff] [blame] | 3 | |
Wenzel Jakob | 8cb6cb3 | 2016-04-17 20:21:41 +0200 | [diff] [blame] | 4 | Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> |
Wenzel Jakob | 1c329aa | 2016-04-13 02:37:36 +0200 | [diff] [blame] | 5 | |
| 6 | All rights reserved. Use of this source code is governed by a |
| 7 | BSD-style license that can be found in the LICENSE file. |
| 8 | */ |
| 9 | |
Dean Moldovan | a0c1ccf | 2016-08-12 13:50:00 +0200 | [diff] [blame] | 10 | #include "pybind11_tests.h" |
Wenzel Jakob | 1c329aa | 2016-04-13 02:37:36 +0200 | [diff] [blame] | 11 | |
| 12 | class Pickleable { |
| 13 | public: |
| 14 | Pickleable(const std::string &value) : m_value(value) { } |
| 15 | const std::string &value() const { return m_value; } |
| 16 | |
| 17 | void setExtra1(int extra1) { m_extra1 = extra1; } |
| 18 | void setExtra2(int extra2) { m_extra2 = extra2; } |
| 19 | int extra1() const { return m_extra1; } |
| 20 | int extra2() const { return m_extra2; } |
| 21 | private: |
| 22 | std::string m_value; |
| 23 | int m_extra1 = 0; |
| 24 | int m_extra2 = 0; |
| 25 | }; |
| 26 | |
Dean Moldovan | 6fccf69 | 2016-10-11 01:12:48 +0200 | [diff] [blame] | 27 | class PickleableWithDict { |
| 28 | public: |
| 29 | PickleableWithDict(const std::string &value) : value(value) { } |
| 30 | |
| 31 | std::string value; |
| 32 | int extra; |
| 33 | }; |
| 34 | |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 35 | test_initializer pickling([](py::module &m) { |
Wenzel Jakob | 1c329aa | 2016-04-13 02:37:36 +0200 | [diff] [blame] | 36 | py::class_<Pickleable>(m, "Pickleable") |
| 37 | .def(py::init<std::string>()) |
| 38 | .def("value", &Pickleable::value) |
| 39 | .def("extra1", &Pickleable::extra1) |
| 40 | .def("extra2", &Pickleable::extra2) |
| 41 | .def("setExtra1", &Pickleable::setExtra1) |
| 42 | .def("setExtra2", &Pickleable::setExtra2) |
| 43 | // For details on the methods below, refer to |
| 44 | // http://docs.python.org/3/library/pickle.html#pickling-class-instances |
| 45 | .def("__getstate__", [](const Pickleable &p) { |
| 46 | /* Return a tuple that fully encodes the state of the object */ |
| 47 | return py::make_tuple(p.value(), p.extra1(), p.extra2()); |
| 48 | }) |
| 49 | .def("__setstate__", [](Pickleable &p, py::tuple t) { |
| 50 | if (t.size() != 3) |
| 51 | throw std::runtime_error("Invalid state!"); |
| 52 | /* Invoke the constructor (need to use in-place version) */ |
| 53 | new (&p) Pickleable(t[0].cast<std::string>()); |
| 54 | |
| 55 | /* Assign any additional state */ |
| 56 | p.setExtra1(t[1].cast<int>()); |
| 57 | p.setExtra2(t[2].cast<int>()); |
| 58 | }); |
Dean Moldovan | 6fccf69 | 2016-10-11 01:12:48 +0200 | [diff] [blame] | 59 | |
Wenzel Jakob | 64cb699 | 2016-12-26 13:12:10 +0100 | [diff] [blame] | 60 | #if !defined(PYPY_VERSION) |
Dean Moldovan | 6fccf69 | 2016-10-11 01:12:48 +0200 | [diff] [blame] | 61 | py::class_<PickleableWithDict>(m, "PickleableWithDict", py::dynamic_attr()) |
| 62 | .def(py::init<std::string>()) |
| 63 | .def_readwrite("value", &PickleableWithDict::value) |
| 64 | .def_readwrite("extra", &PickleableWithDict::extra) |
| 65 | .def("__getstate__", [](py::object self) { |
| 66 | /* Also include __dict__ in state */ |
| 67 | return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__")); |
| 68 | }) |
| 69 | .def("__setstate__", [](py::object self, py::tuple t) { |
| 70 | if (t.size() != 3) |
| 71 | throw std::runtime_error("Invalid state!"); |
| 72 | /* Cast and construct */ |
| 73 | auto& p = self.cast<PickleableWithDict&>(); |
Wenzel Jakob | cd7eacc | 2017-01-04 09:01:59 -0500 | [diff] [blame] | 74 | new (&p) PickleableWithDict(t[0].cast<std::string>()); |
Dean Moldovan | 6fccf69 | 2016-10-11 01:12:48 +0200 | [diff] [blame] | 75 | |
| 76 | /* Assign C++ state */ |
| 77 | p.extra = t[1].cast<int>(); |
| 78 | |
| 79 | /* Assign Python state */ |
| 80 | self.attr("__dict__") = t[2]; |
| 81 | }); |
Wenzel Jakob | 64cb699 | 2016-12-26 13:12:10 +0100 | [diff] [blame] | 82 | #endif |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 83 | }); |