test_pickling.cpp revision 12037:d28054ac6ec9 
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
12class Pickleable {
13public:
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; }
21private:
22    std::string m_value;
23    int m_extra1 = 0;
24    int m_extra2 = 0;
25};
26
27class PickleableWithDict {
28public:
29    PickleableWithDict(const std::string &value) : value(value) { }
30
31    std::string value;
32    int extra;
33};
34
35test_initializer pickling([](py::module &m) {
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        });
59
60#if !defined(PYPY_VERSION)
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&>();
74            new (&p) PickleableWithDict(t[0].cast<std::string>());
75
76            /* Assign C++ state */
77            p.extra = t[1].cast<int>();
78
79            /* Assign Python state */
80            self.attr("__dict__") = t[2];
81        });
82#endif
83});
84