pybind11-2.10.3/tests/test_pickling.cpp

0001 /*
0002     tests/test_pickling.cpp -- pickle support
0003
0004     Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
0005     Copyright (c) 2021 The Pybind Development Team.
0006
0007     All rights reserved. Use of this source code is governed by a
0008     BSD-style license that can be found in the LICENSE file.
0009 */
0010
0011 #include "pybind11_tests.h"
0012
0013 #include <memory>
0014 #include <stdexcept>
0015 #include <utility>
0016
0017 namespace exercise_trampoline {
0018
0019 struct SimpleBase {
0020     int num = 0;
0021     virtual ~SimpleBase() = default;
0022
0023     // For compatibility with old clang versions:
0024     SimpleBase() = default;
0025     SimpleBase(const SimpleBase &) = default;
0026 };
0027
0028 struct SimpleBaseTrampoline : SimpleBase {};
0029
0030 struct SimpleCppDerived : SimpleBase {};
0031
0032 void wrap(py::module m) {
0033     py::class_<SimpleBase, SimpleBaseTrampoline>(m, "SimpleBase")
0034         .def(py::init<>())
0035         .def_readwrite("num", &SimpleBase::num)
0036         .def(py::pickle(
0037             [](const py::object &self) {
0038                 py::dict d;
0039                 if (py::hasattr(self, "__dict__")) {
0040                     d = self.attr("__dict__");
0041                 }
0042                 return py::make_tuple(self.attr("num"), d);
0043             },
0044             [](const py::tuple &t) {
0045                 if (t.size() != 2) {
0046                     throw std::runtime_error("Invalid state!");
0047                 }
0048                 auto cpp_state = std::unique_ptr<SimpleBase>(new SimpleBaseTrampoline);
0049                 cpp_state->num = t[0].cast<int>();
0050                 auto py_state = t[1].cast<py::dict>();
0051                 return std::make_pair(std::move(cpp_state), py_state);
0052             }));
0053
0054     m.def("make_SimpleCppDerivedAsBase",
0055           []() { return std::unique_ptr<SimpleBase>(new SimpleCppDerived); });
0056     m.def("check_dynamic_cast_SimpleCppDerived", [](const SimpleBase *base_ptr) {
0057         return dynamic_cast<const SimpleCppDerived *>(base_ptr) != nullptr;
0058     });
0059 }
0060
0061 } // namespace exercise_trampoline
0062
0063 TEST_SUBMODULE(pickling, m) {
0064     m.def("simple_callable", []() { return 20220426; });
0065
0066     // test_roundtrip
0067     class Pickleable {
0068     public:
0069         explicit Pickleable(const std::string &value) : m_value(value) {}
0070         const std::string &value() const { return m_value; }
0071
0072         void setExtra1(int extra1) { m_extra1 = extra1; }
0073         void setExtra2(int extra2) { m_extra2 = extra2; }
0074         int extra1() const { return m_extra1; }
0075         int extra2() const { return m_extra2; }
0076
0077     private:
0078         std::string m_value;
0079         int m_extra1 = 0;
0080         int m_extra2 = 0;
0081     };
0082
0083     class PickleableNew : public Pickleable {
0084     public:
0085         using Pickleable::Pickleable;
0086     };
0087
0088     py::class_<Pickleable> pyPickleable(m, "Pickleable");
0089     pyPickleable.def(py::init<std::string>())
0090         .def("value", &Pickleable::value)
0091         .def("extra1", &Pickleable::extra1)
0092         .def("extra2", &Pickleable::extra2)
0093         .def("setExtra1", &Pickleable::setExtra1)
0094         .def("setExtra2", &Pickleable::setExtra2)
0095         // For details on the methods below, refer to
0096         // http://docs.python.org/3/library/pickle.html#pickling-class-instances
0097         .def("__getstate__", [](const Pickleable &p) {
0098             /* Return a tuple that fully encodes the state of the object */
0099             return py::make_tuple(p.value(), p.extra1(), p.extra2());
0100         });
0101     ignoreOldStyleInitWarnings([&pyPickleable]() {
0102         pyPickleable.def("__setstate__", [](Pickleable &p, const py::tuple &t) {
0103             if (t.size() != 3) {
0104                 throw std::runtime_error("Invalid state!");
0105             }
0106             /* Invoke the constructor (need to use in-place version) */
0107             new (&p) Pickleable(t[0].cast<std::string>());
0108
0109             /* Assign any additional state */
0110             p.setExtra1(t[1].cast<int>());
0111             p.setExtra2(t[2].cast<int>());
0112         });
0113     });
0114
0115     py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
0116         .def(py::init<std::string>())
0117         .def(py::pickle(
0118             [](const PickleableNew &p) {
0119                 return py::make_tuple(p.value(), p.extra1(), p.extra2());
0120             },
0121             [](const py::tuple &t) {
0122                 if (t.size() != 3) {
0123                     throw std::runtime_error("Invalid state!");
0124                 }
0125                 auto p = PickleableNew(t[0].cast<std::string>());
0126
0127                 p.setExtra1(t[1].cast<int>());
0128                 p.setExtra2(t[2].cast<int>());
0129                 return p;
0130             }));
0131
0132 #if !defined(PYPY_VERSION)
0133     // test_roundtrip_with_dict
0134     class PickleableWithDict {
0135     public:
0136         explicit PickleableWithDict(const std::string &value) : value(value) {}
0137
0138         std::string value;
0139         int extra;
0140     };
0141
0142     class PickleableWithDictNew : public PickleableWithDict {
0143     public:
0144         using PickleableWithDict::PickleableWithDict;
0145     };
0146
0147     py::class_<PickleableWithDict> pyPickleableWithDict(
0148         m, "PickleableWithDict", py::dynamic_attr());
0149     pyPickleableWithDict.def(py::init<std::string>())
0150         .def_readwrite("value", &PickleableWithDict::value)
0151         .def_readwrite("extra", &PickleableWithDict::extra)
0152         .def("__getstate__", [](const py::object &self) {
0153             /* Also include __dict__ in state */
0154             return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
0155         });
0156     ignoreOldStyleInitWarnings([&pyPickleableWithDict]() {
0157         pyPickleableWithDict.def("__setstate__", [](const py::object &self, const py::tuple &t) {
0158             if (t.size() != 3) {
0159                 throw std::runtime_error("Invalid state!");
0160             }
0161             /* Cast and construct */
0162             auto &p = self.cast<PickleableWithDict &>();
0163             new (&p) PickleableWithDict(t[0].cast<std::string>());
0164
0165             /* Assign C++ state */
0166             p.extra = t[1].cast<int>();
0167
0168             /* Assign Python state */
0169             self.attr("__dict__") = t[2];
0170         });
0171     });
0172
0173     py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
0174         .def(py::init<std::string>())
0175         .def(py::pickle(
0176             [](const py::object &self) {
0177                 return py::make_tuple(
0178                     self.attr("value"), self.attr("extra"), self.attr("__dict__"));
0179             },
0180             [](const py::tuple &t) {
0181                 if (t.size() != 3) {
0182                     throw std::runtime_error("Invalid state!");
0183                 }
0184
0185                 auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
0186                 cpp_state.extra = t[1].cast<int>();
0187
0188                 auto py_state = t[2].cast<py::dict>();
0189                 return std::make_pair(cpp_state, py_state);
0190             }));
0191 #endif
0192
0193     exercise_trampoline::wrap(m);
0194 }