From b2462c1d8c97cda460a9ecb21b45332171ac2d78 Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Fri, 1 Jan 2021 17:42:30 +0100
Subject: [PATCH] express: Support pickling PointerToArray objects

---
 panda/src/express/pointerToArray.h     |  4 ++++
 panda/src/express/pointerToArray_ext.I | 28 ++++++++++++++++++++++++++
 panda/src/express/pointerToArray_ext.h |  4 ++++
 tests/express/test_pointertoarray.py   | 23 +++++++++++++++++++++
 4 files changed, 59 insertions(+)
 create mode 100644 tests/express/test_pointertoarray.py

diff --git a/panda/src/express/pointerToArray.h b/panda/src/express/pointerToArray.h
index 7ff07a5ef1..24a5322bc7 100644
--- a/panda/src/express/pointerToArray.h
+++ b/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
diff --git a/panda/src/express/pointerToArray_ext.I b/panda/src/express/pointerToArray_ext.I
index 5bcd1b77e9..435ebb8659 100644
--- a/panda/src/express/pointerToArray_ext.I
+++ b/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.
diff --git a/panda/src/express/pointerToArray_ext.h b/panda/src/express/pointerToArray_ext.h
index f17d1b634a..ce6104791b 100644
--- a/panda/src/express/pointerToArray_ext.h
+++ b/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;
 };
diff --git a/tests/express/test_pointertoarray.py b/tests/express/test_pointertoarray.py
new file mode 100644
index 0000000000..fe76c94128
--- /dev/null
+++ b/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()