express: Support pickling PointerToArray objects

panda/src/express/pointerToArray.h

@@ -117,6 +117,8 @@ PUBLISHED:
   INLINE size_t count(const Element &) const;
 
 #ifdef HAVE_PYTHON
+  EXTENSION(PyObject *__reduce__(PyObject *self) const);
+
   EXTENSION(int __getbuffer__(PyObject *self, Py_buffer *view, int flags));
   EXTENSION(void __releasebuffer__(PyObject *self, Py_buffer *view) const);
 #endif
@@ -273,6 +275,8 @@ PUBLISHED:
   INLINE size_t count(const Element &) const;
 
 #ifdef HAVE_PYTHON
+  EXTENSION(PyObject *__reduce__(PyObject *self) const);
+
   EXTENSION(int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const);
   EXTENSION(void __releasebuffer__(PyObject *self, Py_buffer *view) const);
 #endif

panda/src/express/pointerToArray_ext.I

@@ -259,6 +259,25 @@ get_subdata(size_t n, size_t count) const {
 #endif
 }
 
+/**
+ * Implements pickle support.
+ */
+template<class Element>
+INLINE PyObject *Extension<PointerToArray<Element> >::
+__reduce__(PyObject *self) const {
+  // This preserves the distinction between a null vs. an empty PTA, though I'm
+  // not sure that this distinction matters to anyone.
+  if (this->_this->is_null()) {
+    return Py_BuildValue("O()", Py_TYPE(self));
+  }
+  else if (this->_this->empty()) {
+    return Py_BuildValue("O(())", Py_TYPE(self));
+  }
+  else {
+    return Py_BuildValue("O(N)", Py_TYPE(self), get_data());
+  }
+}
+
 /**
  * Same as get_element(), this returns the nth element of the array.
  */
@@ -304,6 +323,15 @@ get_subdata(size_t n, size_t count) const {
 #endif
 }
 
+/**
+ * Implements pickle support.
+ */
+template<class Element>
+INLINE PyObject *Extension<ConstPointerToArray<Element> >::
+__reduce__(PyObject *self) const {
+  return Py_BuildValue("O(N)", Py_TYPE(self), get_data());
+}
+
 /**
  * This is used to implement the buffer protocol, in order to allow efficient
  * access to the array data through a Python multiview object.

panda/src/express/pointerToArray_ext.h

@@ -40,6 +40,8 @@ public:
   INLINE void set_data(PyObject *data);
   INLINE PyObject *get_subdata(size_t n, size_t count) const;
 
+  INLINE PyObject *__reduce__(PyObject *self) const;
+
   INLINE int __getbuffer__(PyObject *self, Py_buffer *view, int flags);
   INLINE void __releasebuffer__(PyObject *self, Py_buffer *view) const;
 };
@@ -75,6 +77,8 @@ public:
   INLINE PyObject *get_data() const;
   INLINE PyObject *get_subdata(size_t n, size_t count) const;
 
+  INLINE PyObject *__reduce__(PyObject *self) const;
+
   INLINE int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const;
   INLINE void __releasebuffer__(PyObject *self, Py_buffer *view) const;
 };

tests/express/test_pointertoarray.py

@@ -0,0 +1,23 @@
+def test_pta_float_pickle():
+    from panda3d.core import PTA_float
+    from direct.stdpy.pickle import dumps, loads, HIGHEST_PROTOCOL
+
+    null_pta = PTA_float()
+
+    empty_pta = PTA_float([])
+
+    data_pta = PTA_float([1.0, 2.0, 3.0])
+    data = data_pta.get_data()
+
+    for proto in range(1, HIGHEST_PROTOCOL + 1):
+        null_pta2 = loads(dumps(null_pta, proto))
+        assert null_pta2.is_null()
+        assert len(null_pta2) == 0
+
+        empty_pta2 = loads(dumps(empty_pta, proto))
+        assert not empty_pta2.is_null()
+        assert len(empty_pta2) == 0
+
+        data_pta2 = loads(dumps(data_pta, proto))
+        assert tuple(data_pta2) == (1.0, 2.0, 3.0)
+        assert data_pta2.get_data() == data_pta.get_data()