putil: add pickling support to Datagram class

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;
 

panda/src/express/datagram_ext.I

@@ -33,3 +33,29 @@ PyObject *Extension<Datagram>::
 __bytes__() const {
   return get_message();
 }
+
+/**
+ * This special Python method is implemented to provide support for the pickle
+ * module.
+ */
+INLINE PyObject *Extension<Datagram>::
+__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;
+}

panda/src/express/datagram_ext.h

@@ -31,6 +31,7 @@ class Extension<Datagram> : public ExtensionBase<Datagram> {
 public:
   INLINE PyObject *get_message() const;
   INLINE PyObject *__bytes__() const;
+  INLINE PyObject *__reduce__() const;
 };
 
 #include "datagram_ext.I"

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."""