1 /*
2 tests/test_pickling.cpp -- pickle support
3
4 Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
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
10 #include "pybind11_tests.h"
11
TEST_SUBMODULE(pickling,m)12 TEST_SUBMODULE(pickling, m) {
13 // test_roundtrip
14 class Pickleable {
15 public:
16 Pickleable(const std::string &value) : m_value(value) { }
17 const std::string &value() const { return m_value; }
18
19 void setExtra1(int extra1) { m_extra1 = extra1; }
20 void setExtra2(int extra2) { m_extra2 = extra2; }
21 int extra1() const { return m_extra1; }
22 int extra2() const { return m_extra2; }
23 private:
24 std::string m_value;
25 int m_extra1 = 0;
26 int m_extra2 = 0;
27 };
28
29 class PickleableNew : public Pickleable {
30 public:
31 using Pickleable::Pickleable;
32 };
33
34 py::class_<Pickleable> pyPickleable(m, "Pickleable");
35 pyPickleable
36 .def(py::init<std::string>())
37 .def("value", &Pickleable::value)
38 .def("extra1", &Pickleable::extra1)
39 .def("extra2", &Pickleable::extra2)
40 .def("setExtra1", &Pickleable::setExtra1)
41 .def("setExtra2", &Pickleable::setExtra2)
42 // For details on the methods below, refer to
43 // http://docs.python.org/3/library/pickle.html#pickling-class-instances
44 .def("__getstate__", [](const Pickleable &p) {
45 /* Return a tuple that fully encodes the state of the object */
46 return py::make_tuple(p.value(), p.extra1(), p.extra2());
47 });
48 ignoreOldStyleInitWarnings([&pyPickleable]() {
49 pyPickleable
50 .def("__setstate__", [](Pickleable &p, py::tuple t) {
51 if (t.size() != 3)
52 throw std::runtime_error("Invalid state!");
53 /* Invoke the constructor (need to use in-place version) */
54 new (&p) Pickleable(t[0].cast<std::string>());
55
56 /* Assign any additional state */
57 p.setExtra1(t[1].cast<int>());
58 p.setExtra2(t[2].cast<int>());
59 });
60 });
61
62 py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
63 .def(py::init<std::string>())
64 .def(py::pickle(
65 [](const PickleableNew &p) {
66 return py::make_tuple(p.value(), p.extra1(), p.extra2());
67 },
68 [](py::tuple t) {
69 if (t.size() != 3)
70 throw std::runtime_error("Invalid state!");
71 auto p = PickleableNew(t[0].cast<std::string>());
72
73 p.setExtra1(t[1].cast<int>());
74 p.setExtra2(t[2].cast<int>());
75 return p;
76 }
77 ));
78
79 #if !defined(PYPY_VERSION)
80 // test_roundtrip_with_dict
81 class PickleableWithDict {
82 public:
83 PickleableWithDict(const std::string &value) : value(value) { }
84
85 std::string value;
86 int extra;
87 };
88
89 class PickleableWithDictNew : public PickleableWithDict {
90 public:
91 using PickleableWithDict::PickleableWithDict;
92 };
93
94 py::class_<PickleableWithDict> pyPickleableWithDict(m, "PickleableWithDict", py::dynamic_attr());
95 pyPickleableWithDict
96 .def(py::init<std::string>())
97 .def_readwrite("value", &PickleableWithDict::value)
98 .def_readwrite("extra", &PickleableWithDict::extra)
99 .def("__getstate__", [](py::object self) {
100 /* Also include __dict__ in state */
101 return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
102 });
103 ignoreOldStyleInitWarnings([&pyPickleableWithDict]() {
104 pyPickleableWithDict
105 .def("__setstate__", [](py::object self, py::tuple t) {
106 if (t.size() != 3)
107 throw std::runtime_error("Invalid state!");
108 /* Cast and construct */
109 auto& p = self.cast<PickleableWithDict&>();
110 new (&p) PickleableWithDict(t[0].cast<std::string>());
111
112 /* Assign C++ state */
113 p.extra = t[1].cast<int>();
114
115 /* Assign Python state */
116 self.attr("__dict__") = t[2];
117 });
118 });
119
120 py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
121 .def(py::init<std::string>())
122 .def(py::pickle(
123 [](py::object self) {
124 return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
125 },
126 [](const py::tuple &t) {
127 if (t.size() != 3)
128 throw std::runtime_error("Invalid state!");
129
130 auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
131 cpp_state.extra = t[1].cast<int>();
132
133 auto py_state = t[2].cast<py::dict>();
134 return std::make_pair(cpp_state, py_state);
135 }
136 ));
137 #endif
138 }
139