summaryrefslogtreecommitdiff
path: root/ext/pybind11/tests/test_pickling.cpp
blob: 9dc63bda3b5949032fbcd30e7aa4e7db2072dcff (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
/*
    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
}