Add pickle support to a variety of classes

dtool/src/dtoolbase/typeHandle.h

@@ -137,6 +137,9 @@ PUBLISHED:
   MAKE_SEQ_PROPERTY(parent_classes, get_num_parent_classes, get_parent_class);
   MAKE_SEQ_PROPERTY(child_classes, get_num_child_classes, get_child_class);
 
+  EXTENSION(PyObject *__reduce__() const);
+  EXTENSION(void __setstate__(PyObject *));
+
 public:
 #ifdef HAVE_PYTHON
   PyObject *get_python_type() const;

dtool/src/dtoolbase/typeHandle_ext.cxx

@@ -32,4 +32,55 @@ make(PyTypeObject *tp) {
   return dtool_tp->_type;
 }
 
+/**
+ * Implements pickle support.
+ */
+PyObject *Extension<TypeHandle>::
+__reduce__() const {
+  extern struct Dtool_PyTypedObject Dtool_TypeHandle;
+  extern struct Dtool_PyTypedObject Dtool_TypeRegistry;
+
+  if (!*_this) {
+    PyObject *func = PyObject_GetAttrString((PyObject *)&Dtool_TypeHandle, "none");
+    return Py_BuildValue("N()", func);
+  }
+
+  // If we have a Python binding registered for it, that's the preferred method,
+  // since it ensures that the appropriate module gets loaded by pickle.
+  PyObject *py_type = _this->get_python_type();
+  if (py_type != nullptr && *_this == ((Dtool_PyTypedObject *)py_type)->_type) {
+    PyObject *func = PyObject_GetAttrString((PyObject *)&Dtool_TypeHandle, "make");
+    return Py_BuildValue("N(O)", func, py_type);
+  }
+
+  // Fall back to the __setstate__ mechanism.
+  std::string name = _this->get_name();
+  Py_ssize_t num_parents = _this->get_num_parent_classes();
+  PyObject *parents = PyTuple_New(num_parents);
+  for (Py_ssize_t i = 0; i < num_parents; ++i) {
+    PyObject *parent = DTool_CreatePyInstance(new TypeHandle(_this->get_parent_class(i)), Dtool_TypeHandle, true, false);
+    PyTuple_SET_ITEM(parents, i, parent);
+  }
+  return Py_BuildValue("O()(s#N)", (PyObject *)&Dtool_TypeHandle, name.c_str(), name.size(), parents);
+}
+
+/**
+ * Implements pickle support.
+ */
+void Extension<TypeHandle>::
+__setstate__(PyObject *state) {
+  Py_ssize_t len;
+  const char *name_str = PyUnicode_AsUTF8AndSize(PyTuple_GET_ITEM(state, 0), &len);
+  PyObject *parents = PyTuple_GET_ITEM(state, 1);
+
+  TypeRegistry *type_registry = TypeRegistry::ptr();
+  *_this = type_registry->register_dynamic_type(std::string(name_str, len));
+
+  Py_ssize_t num_parents = PyTuple_GET_SIZE(parents);
+  for (Py_ssize_t i = 0; i < num_parents; ++i) {
+    TypeHandle *parent = (TypeHandle *)DtoolInstance_VOID_PTR(PyTuple_GET_ITEM(parents, i));
+    type_registry->record_derivation(*_this, *parent);
+  }
+}
+
 #endif

dtool/src/dtoolbase/typeHandle_ext.h

@@ -30,6 +30,9 @@ template<>
 class Extension<TypeHandle> : public ExtensionBase<TypeHandle> {
 public:
   static TypeHandle make(PyTypeObject *tp);
+
+  PyObject *__reduce__() const;
+  void __setstate__(PyObject *);
 };
 
 #endif  // HAVE_PYTHON

panda/src/display/CMakeLists.txt

@@ -67,6 +67,8 @@ set(P3DISPLAY_SOURCES
 )
 
 set(P3DISPLAY_IGATEEXT
+  graphicsPipeSelection_ext.cxx
+  graphicsPipeSelection_ext.h
   graphicsStateGuardian_ext.cxx
   graphicsStateGuardian_ext.h
   graphicsWindow_ext.cxx

panda/src/display/graphicsPipeSelection.h

@@ -52,6 +52,8 @@ PUBLISHED:
 
   INLINE static GraphicsPipeSelection *get_global_ptr();
 
+  EXTENSION(PyObject *__reduce__() const);
+
 public:
   typedef PT(GraphicsPipe) PipeConstructorFunc();
   bool add_pipe_type(TypeHandle type, PipeConstructorFunc *func);

panda/src/display/graphicsPipeSelection_ext.cxx

@@ -0,0 +1,31 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file graphicsPipeSelection_ext.cxx
+ * @author rdb
+ * @date 2021-12-10
+ */
+
+#include "graphicsPipeSelection_ext.h"
+
+#ifdef HAVE_PYTHON
+
+#include "pythonLoaderFileType.h"
+
+extern struct Dtool_PyTypedObject Dtool_GraphicsPipeSelection;
+
+/**
+ * Implements pickle support.
+ */
+PyObject *Extension<GraphicsPipeSelection>::
+__reduce__() const {
+  PyObject *func = PyObject_GetAttrString((PyObject *)&Dtool_GraphicsPipeSelection, "get_global_ptr");
+  return Py_BuildValue("N()", func);
+}
+
+#endif

panda/src/display/graphicsPipeSelection_ext.h

@@ -0,0 +1,37 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file graphicsPipeSelection_ext.h
+ * @author rdb
+ * @date 2021-12-10
+ */
+
+#ifndef GRAPHICSPIPESELECTION_EXT_H
+#define GRAPHICSPIPESELECTION_EXT_H
+
+#include "pandabase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "graphicsPipeSelection.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for GraphicsPipeSelection, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<>
+class Extension<GraphicsPipeSelection> : public ExtensionBase<GraphicsPipeSelection> {
+public:
+  PyObject *__reduce__() const;
+};
+
+#endif  // HAVE_PYTHON
+
+#endif  // GRAPHICSPIPESELECTION_EXT_H

panda/src/display/p3display_ext_composite.cxx

@@ -1,3 +1,4 @@
+#include "graphicsPipeSelection_ext.cxx"
 #include "graphicsStateGuardian_ext.cxx"
 #include "graphicsWindow_ext.cxx"
 #include "pythonGraphicsWindowProc.cxx"

panda/src/display/windowProperties.h

@@ -205,6 +205,9 @@ PUBLISHED:
   MAKE_PROPERTY2(parent_window, has_parent_window, get_parent_window,
                                 set_parent_window, clear_parent_window);
 
+  EXTENSION(PyObject *__getstate__(PyObject *self) const);
+  EXTENSION(void __setstate__(PyObject *self, PyObject *state));
+
   void add_properties(const WindowProperties &other);
 
   void output(std::ostream &out) const;

panda/src/display/windowProperties_ext.cxx

@@ -50,27 +50,63 @@ __init__(PyObject *self, PyObject *args, PyObject *kwds) {
   // Now iterate over the keyword arguments, which define the default values
   // for the different properties.
   if (kwds != nullptr) {
-    PyTypeObject *type = Py_TYPE(self);
-    PyObject *key, *value;
-    Py_ssize_t pos = 0;
-
-    while (PyDict_Next(kwds, &pos, &key, &value)) {
-      // Look for a writable property on the type by this name.
-      PyObject *descr = _PyType_Lookup(type, key);
-
-      if (descr != nullptr && Py_TYPE(descr)->tp_descr_set != nullptr) {
-        if (Py_TYPE(descr)->tp_descr_set(descr, self, value) < 0) {
-          return;
-        }
-      } else {
-        PyObject *key_repr = PyObject_Repr(key);
-        PyErr_Format(PyExc_TypeError,
-                     "%.100s is an invalid keyword argument for WindowProperties()",
-                     PyUnicode_AsUTF8(key_repr)
-                    );
-        Py_DECREF(key_repr);
+    __setstate__(self, kwds);
+  }
+}
+
+/**
+ * Returns the properties as a dictionary.
+ */
+PyObject *Extension<WindowProperties>::
+__getstate__(PyObject *self) const {
+  static const char *props[] = {"origin", "size", "title", "undecorated", "fixed_size", "fullscreen", "foreground", "minimized", "maximized", "raw_mice", "open", "cursor_hidden", "icon_filename", "cursor_filename", "z_order", "mouse_mode", "parent_window", nullptr};
+
+  PyTypeObject *type = Py_TYPE(self);
+  PyObject *state = PyDict_New();
+
+  for (size_t i = 0; props[i] != nullptr; ++i) {
+    PyObject *key = PyUnicode_FromString(props[i]);
+    PyObject *descr = _PyType_Lookup(type, key);
+
+    if (descr != nullptr && Py_TYPE(descr)->tp_descr_get != nullptr) {
+      PyObject *value = Py_TYPE(descr)->tp_descr_get(descr, self, (PyObject *)type);
+      nassertr(value != nullptr, nullptr);
+      if (value != Py_None) {
+        PyDict_SetItem(state, key, value);
+      }
+      Py_DECREF(value);
+    }
+    Py_DECREF(key);
+  }
+
+  return state;
+}
+
+/**
+ *
+ */
+void Extension<WindowProperties>::
+__setstate__(PyObject *self, PyObject *props) {
+  PyTypeObject *type = Py_TYPE(self);
+  PyObject *key, *value;
+  Py_ssize_t pos = 0;
+
+  while (PyDict_Next(props, &pos, &key, &value)) {
+    // Look for a writable property on the type by this name.
+    PyObject *descr = _PyType_Lookup(type, key);
+
+    if (descr != nullptr && Py_TYPE(descr)->tp_descr_set != nullptr) {
+      if (Py_TYPE(descr)->tp_descr_set(descr, self, value) < 0) {
         return;
       }
+    } else {
+      PyObject *key_repr = PyObject_Repr(key);
+      PyErr_Format(PyExc_TypeError,
+                   "%.100s is an invalid keyword argument for WindowProperties()",
+                   PyUnicode_AsUTF8(key_repr)
+                  );
+      Py_DECREF(key_repr);
+      return;
     }
   }
 }

panda/src/display/windowProperties_ext.h

@@ -30,6 +30,9 @@ template<>
 class Extension<WindowProperties> : public ExtensionBase<WindowProperties> {
 public:
   void __init__(PyObject *self, PyObject *args, PyObject *kwds);
+
+  PyObject *__getstate__(PyObject *self) const;
+  void __setstate__(PyObject *self, PyObject *state);
 };
 
 #endif  // HAVE_PYTHON

panda/src/express/ramfile.h

@@ -36,6 +36,9 @@ PUBLISHED:
   INLINE size_t get_data_size() const;
   INLINE void clear();
 
+  EXTENSION(PyObject *__getstate__() const);
+  EXTENSION(void __setstate__(PyObject *state));
+
 public:
   std::string read(size_t length);
   std::string readline();

panda/src/express/ramfile_ext.cxx

@@ -74,4 +74,28 @@ get_data() const {
   return PyBytes_FromStringAndSize(_this->_data.data(), _this->_data.size());
 }
 
+/**
+ *
+ */
+PyObject *Extension<Ramfile>::
+__getstate__() const {
+  PyObject *state = PyTuple_New(2);
+  PyTuple_SET_ITEM(state, 0, get_data());
+  PyTuple_SET_ITEM(state, 1, PyLong_FromSize_t(_this->tell()));
+  return state;
+}
+
+/**
+ *
+ */
+void Extension<Ramfile>::
+__setstate__(PyObject *state) {
+  char *str;
+  Py_ssize_t len;
+  if (PyBytes_AsStringAndSize(PyTuple_GET_ITEM(state, 0), &str, &len) >= 0) {
+    _this->_data = std::string(str, len);
+  }
+  _this->seek(PyLong_AsSize_t(PyTuple_GET_ITEM(state, 1)));
+}
+
 #endif

panda/src/express/ramfile_ext.h

@@ -34,6 +34,9 @@ public:
   PyObject *readlines();
 
   PyObject *get_data() const;
+
+  PyObject *__getstate__() const;
+  void __setstate__(PyObject *);
 };
 
 #endif  // HAVE_PYTHON

panda/src/putil/pythonCallbackObject.cxx

@@ -29,6 +29,7 @@ ConfigureFn(config_pythonCallbackObject) {
 
 #ifndef CPPPARSER
 extern struct Dtool_PyTypedObject Dtool_TypedObject;
+extern struct Dtool_PyTypedObject Dtool_PythonCallbackObject;
 #endif
 
 /**
@@ -84,6 +85,14 @@ get_function() {
   return _function;
 }
 
+/**
+ * Implements pickle support.
+ */
+PyObject *PythonCallbackObject::
+__reduce__() const {
+  return Py_BuildValue("O(O)", (PyObject *)&Dtool_PythonCallbackObject, _function);
+}
+
 /**
  * This method called when the callback is triggered; it *replaces* the
  * original function.  To continue performing the original function, you must

panda/src/putil/pythonCallbackObject.h

@@ -23,7 +23,7 @@
 
 /**
  * This is a specialization on CallbackObject to allow a callback to directly
- * call an arbitarary Python function.  Powerful!  But use with caution.
+ * call an arbitrary Python function.  Powerful!  But use with caution.
  */
 class PythonCallbackObject : public CallbackObject {
 PUBLISHED:
@@ -34,6 +34,8 @@ PUBLISHED:
   void set_function(PyObject *function);
   PyObject *get_function();
 
+  PyObject *__reduce__() const;
+
   MAKE_PROPERTY(function, get_function, set_function);
 
 public: