You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

195 lines
6.5 KiB

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