test_pickling.cpp revision 11986
1545SN/A/* 28948Sandreas.hansson@arm.com tests/test_pickling.cpp -- pickle support 38948Sandreas.hansson@arm.com 48948Sandreas.hansson@arm.com Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> 58948Sandreas.hansson@arm.com 68948Sandreas.hansson@arm.com All rights reserved. Use of this source code is governed by a 78948Sandreas.hansson@arm.com BSD-style license that can be found in the LICENSE file. 88948Sandreas.hansson@arm.com*/ 98948Sandreas.hansson@arm.com 108948Sandreas.hansson@arm.com#include "pybind11_tests.h" 118948Sandreas.hansson@arm.com 128948Sandreas.hansson@arm.comclass Pickleable { 138948Sandreas.hansson@arm.compublic: 141762SN/A Pickleable(const std::string &value) : m_value(value) { } 15545SN/A const std::string &value() const { return m_value; } 16545SN/A 17545SN/A void setExtra1(int extra1) { m_extra1 = extra1; } 18545SN/A void setExtra2(int extra2) { m_extra2 = extra2; } 19545SN/A int extra1() const { return m_extra1; } 20545SN/A int extra2() const { return m_extra2; } 21545SN/Aprivate: 22545SN/A std::string m_value; 23545SN/A int m_extra1 = 0; 24545SN/A int m_extra2 = 0; 25545SN/A}; 26545SN/A 27545SN/Aclass PickleableWithDict { 28545SN/Apublic: 29545SN/A PickleableWithDict(const std::string &value) : value(value) { } 30545SN/A 31545SN/A std::string value; 32545SN/A int extra; 33545SN/A}; 34545SN/A 35545SN/Atest_initializer pickling([](py::module &m) { 36545SN/A py::class_<Pickleable>(m, "Pickleable") 37545SN/A .def(py::init<std::string>()) 38545SN/A .def("value", &Pickleable::value) 392665Ssaidi@eecs.umich.edu .def("extra1", &Pickleable::extra1) 402665Ssaidi@eecs.umich.edu .def("extra2", &Pickleable::extra2) 412665Ssaidi@eecs.umich.edu .def("setExtra1", &Pickleable::setExtra1) 42545SN/A .def("setExtra2", &Pickleable::setExtra2) 43545SN/A // For details on the methods below, refer to 441310SN/A // http://docs.python.org/3/library/pickle.html#pickling-class-instances 451310SN/A .def("__getstate__", [](const Pickleable &p) { 46545SN/A /* Return a tuple that fully encodes the state of the object */ 472542SN/A return py::make_tuple(p.value(), p.extra1(), p.extra2()); 483348Sbinkertn@umich.edu }) 494762Snate@binkert.org .def("__setstate__", [](Pickleable &p, py::tuple t) { 504762Snate@binkert.org if (t.size() != 3) 51545SN/A throw std::runtime_error("Invalid state!"); 522384SN/A /* Invoke the constructor (need to use in-place version) */ 532522SN/A new (&p) Pickleable(t[0].cast<std::string>()); 54545SN/A 552489SN/A /* Assign any additional state */ 562489SN/A p.setExtra1(t[1].cast<int>()); 572489SN/A p.setExtra2(t[2].cast<int>()); 582489SN/A }); 598711Sandreas.hansson@arm.com 603090Sstever@eecs.umich.edu py::class_<PickleableWithDict>(m, "PickleableWithDict", py::dynamic_attr()) 613090Sstever@eecs.umich.edu .def(py::init<std::string>()) 622914Ssaidi@eecs.umich.edu .def_readwrite("value", &PickleableWithDict::value) 63545SN/A .def_readwrite("extra", &PickleableWithDict::extra) 64545SN/A .def("__getstate__", [](py::object self) { 652489SN/A /* Also include __dict__ in state */ 662384SN/A return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__")); 672384SN/A }) 683349Sbinkertn@umich.edu .def("__setstate__", [](py::object self, py::tuple t) { 692384SN/A if (t.size() != 3) 709090Sandreas.hansson@arm.com throw std::runtime_error("Invalid state!"); 712384SN/A /* Cast and construct */ 722384SN/A auto& p = self.cast<PickleableWithDict&>(); 733091Sstever@eecs.umich.edu new (&p) Pickleable(t[0].cast<std::string>()); 748914Sandreas.hansson@arm.com 752384SN/A /* Assign C++ state */ 762384SN/A p.extra = t[1].cast<int>(); 772489SN/A 782489SN/A /* Assign Python state */ 792489SN/A self.attr("__dict__") = t[2]; 808711Sandreas.hansson@arm.com }); 812489SN/A}); 822489SN/A