diff --git a/panda/src/express/datagram.h b/panda/src/express/datagram.h index 9611645eed..e0590427f0 100644 --- a/panda/src/express/datagram.h +++ b/panda/src/express/datagram.h @@ -97,6 +97,7 @@ public: PUBLISHED: EXTENSION(INLINE PyObject *get_message() const); EXTENSION(INLINE PyObject *__bytes__() const); + EXTENSION(PyObject *__reduce__() const); INLINE size_t get_length() const; diff --git a/panda/src/express/datagram_ext.I b/panda/src/express/datagram_ext.I index a35c5410b2..9c997e64f6 100644 --- a/panda/src/express/datagram_ext.I +++ b/panda/src/express/datagram_ext.I @@ -33,3 +33,29 @@ PyObject *Extension:: __bytes__() const { return get_message(); } + +/** + * This special Python method is implemented to provide support for the pickle + * module. + */ +INLINE PyObject *Extension:: +__reduce__() const { + // We should return at least a 2-tuple, (Class, (args)): the necessary class + // object whose constructor we should call (e.g. this), and the arguments + // necessary to reconstruct this object. + PyObject *args; + if (_this->get_length() > 0) { + args = PyTuple_New(1); + PyTuple_SET_ITEM(args, 0, get_message()); + } else { + args = PyTuple_New(0); + } + + extern struct Dtool_PyTypedObject Dtool_Datagram; + Py_INCREF((PyObject *)&Dtool_Datagram._PyType); + + PyObject *result = PyTuple_New(2); + PyTuple_SET_ITEM(result, 0, (PyObject *)&Dtool_Datagram._PyType); + PyTuple_SET_ITEM(result, 1, args); + return result; +} diff --git a/panda/src/express/datagram_ext.h b/panda/src/express/datagram_ext.h index 35f2e32bc6..ebaf2dd778 100644 --- a/panda/src/express/datagram_ext.h +++ b/panda/src/express/datagram_ext.h @@ -31,6 +31,7 @@ class Extension : public ExtensionBase { public: INLINE PyObject *get_message() const; INLINE PyObject *__bytes__() const; + INLINE PyObject *__reduce__() const; }; #include "datagram_ext.I" diff --git a/tests/putil/test_datagram.py b/tests/putil/test_datagram.py index 5a654b75ed..ad8ef9857c 100644 --- a/tests/putil/test_datagram.py +++ b/tests/putil/test_datagram.py @@ -101,6 +101,16 @@ def test_datagram_get_message(): assert dg.get_message() == b'abc\x00\xff123' +def test_datagram_pickle(): + import pickle + + dg = core.Datagram() + assert pickle.loads(pickle.dumps(dg, -1)) == dg + + dg = core.Datagram(b'abc\x00') + assert pickle.loads(pickle.dumps(dg, -1)) == dg + + def test_iterator(datagram_small): """This tests Datagram/DatagramIterator, and sort of serves as a self-check of the test fixtures too."""