test_pickling.cpp revision 12391:ceeca8b41e4b 
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
12TEST_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>(m, "Pickleable")
35        .def(py::init<std::string>())
36        .def("value", &Pickleable::value)
37        .def("extra1", &Pickleable::extra1)
38        .def("extra2", &Pickleable::extra2)
39        .def("setExtra1", &Pickleable::setExtra1)
40        .def("setExtra2", &Pickleable::setExtra2)
41        // For details on the methods below, refer to
42        // http://docs.python.org/3/library/pickle.html#pickling-class-instances
43        .def("__getstate__", [](const Pickleable &p) {
44            /* Return a tuple that fully encodes the state of the object */
45            return py::make_tuple(p.value(), p.extra1(), p.extra2());
46        })
47        .def("__setstate__", [](Pickleable &p, py::tuple t) {
48            if (t.size() != 3)
49                throw std::runtime_error("Invalid state!");
50            /* Invoke the constructor (need to use in-place version) */
51            new (&p) Pickleable(t[0].cast<std::string>());
52
53            /* Assign any additional state */
54            p.setExtra1(t[1].cast<int>());
55            p.setExtra2(t[2].cast<int>());
56        });
57
58    py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
59        .def(py::init<std::string>())
60        .def(py::pickle(
61            [](const PickleableNew &p) {
62                return py::make_tuple(p.value(), p.extra1(), p.extra2());
63            },
64            [](py::tuple t) {
65                if (t.size() != 3)
66                    throw std::runtime_error("Invalid state!");
67                auto p = PickleableNew(t[0].cast<std::string>());
68
69                p.setExtra1(t[1].cast<int>());
70                p.setExtra2(t[2].cast<int>());
71                return p;
72            }
73        ));
74
75#if !defined(PYPY_VERSION)
76    // test_roundtrip_with_dict
77    class PickleableWithDict {
78    public:
79        PickleableWithDict(const std::string &value) : value(value) { }
80
81        std::string value;
82        int extra;
83    };
84
85    class PickleableWithDictNew : public PickleableWithDict {
86    public:
87        using PickleableWithDict::PickleableWithDict;
88    };
89
90    py::class_<PickleableWithDict>(m, "PickleableWithDict", py::dynamic_attr())
91        .def(py::init<std::string>())
92        .def_readwrite("value", &PickleableWithDict::value)
93        .def_readwrite("extra", &PickleableWithDict::extra)
94        .def("__getstate__", [](py::object self) {
95            /* Also include __dict__ in state */
96            return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
97        })
98        .def("__setstate__", [](py::object self, py::tuple t) {
99            if (t.size() != 3)
100                throw std::runtime_error("Invalid state!");
101            /* Cast and construct */
102            auto& p = self.cast<PickleableWithDict&>();
103            new (&p) PickleableWithDict(t[0].cast<std::string>());
104
105            /* Assign C++ state */
106            p.extra = t[1].cast<int>();
107
108            /* Assign Python state */
109            self.attr("__dict__") = t[2];
110        });
111
112    py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
113        .def(py::init<std::string>())
114        .def(py::pickle(
115            [](py::object self) {
116                return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
117            },
118            [](const py::tuple &t) {
119                if (t.size() != 3)
120                    throw std::runtime_error("Invalid state!");
121
122                auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
123                cpp_state.extra = t[1].cast<int>();
124
125                auto py_state = t[2].cast<py::dict>();
126                return std::make_pair(cpp_state, py_state);
127            }
128        ));
129#endif
130}
131