test_pickling.cpp 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. /*
  2. tests/test_pickling.cpp -- pickle support
  3. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  4. All rights reserved. Use of this source code is governed by a
  5. BSD-style license that can be found in the LICENSE file.
  6. */
  7. #include "pybind11_tests.h"
  8. TEST_SUBMODULE(pickling, m) {
  9. // test_roundtrip
  10. class Pickleable {
  11. public:
  12. Pickleable(const std::string &value) : m_value(value) { }
  13. const std::string &value() const { return m_value; }
  14. void setExtra1(int extra1) { m_extra1 = extra1; }
  15. void setExtra2(int extra2) { m_extra2 = extra2; }
  16. int extra1() const { return m_extra1; }
  17. int extra2() const { return m_extra2; }
  18. private:
  19. std::string m_value;
  20. int m_extra1 = 0;
  21. int m_extra2 = 0;
  22. };
  23. class PickleableNew : public Pickleable {
  24. public:
  25. using Pickleable::Pickleable;
  26. };
  27. py::class_<Pickleable> pyPickleable(m, "Pickleable");
  28. pyPickleable
  29. .def(py::init<std::string>())
  30. .def("value", &Pickleable::value)
  31. .def("extra1", &Pickleable::extra1)
  32. .def("extra2", &Pickleable::extra2)
  33. .def("setExtra1", &Pickleable::setExtra1)
  34. .def("setExtra2", &Pickleable::setExtra2)
  35. // For details on the methods below, refer to
  36. // http://docs.python.org/3/library/pickle.html#pickling-class-instances
  37. .def("__getstate__", [](const Pickleable &p) {
  38. /* Return a tuple that fully encodes the state of the object */
  39. return py::make_tuple(p.value(), p.extra1(), p.extra2());
  40. });
  41. ignoreOldStyleInitWarnings([&pyPickleable]() {
  42. pyPickleable
  43. .def("__setstate__", [](Pickleable &p, py::tuple t) {
  44. if (t.size() != 3)
  45. throw std::runtime_error("Invalid state!");
  46. /* Invoke the constructor (need to use in-place version) */
  47. new (&p) Pickleable(t[0].cast<std::string>());
  48. /* Assign any additional state */
  49. p.setExtra1(t[1].cast<int>());
  50. p.setExtra2(t[2].cast<int>());
  51. });
  52. });
  53. py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
  54. .def(py::init<std::string>())
  55. .def(py::pickle(
  56. [](const PickleableNew &p) {
  57. return py::make_tuple(p.value(), p.extra1(), p.extra2());
  58. },
  59. [](py::tuple t) {
  60. if (t.size() != 3)
  61. throw std::runtime_error("Invalid state!");
  62. auto p = PickleableNew(t[0].cast<std::string>());
  63. p.setExtra1(t[1].cast<int>());
  64. p.setExtra2(t[2].cast<int>());
  65. return p;
  66. }
  67. ));
  68. #if !defined(PYPY_VERSION)
  69. // test_roundtrip_with_dict
  70. class PickleableWithDict {
  71. public:
  72. PickleableWithDict(const std::string &value) : value(value) { }
  73. std::string value;
  74. int extra;
  75. };
  76. class PickleableWithDictNew : public PickleableWithDict {
  77. public:
  78. using PickleableWithDict::PickleableWithDict;
  79. };
  80. py::class_<PickleableWithDict> pyPickleableWithDict(m, "PickleableWithDict", py::dynamic_attr());
  81. pyPickleableWithDict
  82. .def(py::init<std::string>())
  83. .def_readwrite("value", &PickleableWithDict::value)
  84. .def_readwrite("extra", &PickleableWithDict::extra)
  85. .def("__getstate__", [](py::object self) {
  86. /* Also include __dict__ in state */
  87. return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
  88. });
  89. ignoreOldStyleInitWarnings([&pyPickleableWithDict]() {
  90. pyPickleableWithDict
  91. .def("__setstate__", [](py::object self, py::tuple t) {
  92. if (t.size() != 3)
  93. throw std::runtime_error("Invalid state!");
  94. /* Cast and construct */
  95. auto& p = self.cast<PickleableWithDict&>();
  96. new (&p) PickleableWithDict(t[0].cast<std::string>());
  97. /* Assign C++ state */
  98. p.extra = t[1].cast<int>();
  99. /* Assign Python state */
  100. self.attr("__dict__") = t[2];
  101. });
  102. });
  103. py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
  104. .def(py::init<std::string>())
  105. .def(py::pickle(
  106. [](py::object self) {
  107. return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
  108. },
  109. [](const py::tuple &t) {
  110. if (t.size() != 3)
  111. throw std::runtime_error("Invalid state!");
  112. auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
  113. cpp_state.extra = t[1].cast<int>();
  114. auto py_state = t[2].cast<py::dict>();
  115. return std::make_pair(cpp_state, py_state);
  116. }
  117. ));
  118. #endif
  119. }