From e755f8713088c986bd3dc3f23d8e3f609e3a2bbe Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Fri, 1 Jan 2021 17:08:55 +0100
Subject: [PATCH] pgraph: Add pickling for LoaderFileTypeRegistry

Useful to test that pickling singletons works
---
 panda/src/pgraph/loaderFileTypeRegistry.h       |  2 ++
 panda/src/pgraph/loaderFileTypeRegistry_ext.cxx | 10 ++++++++++
 panda/src/pgraph/loaderFileTypeRegistry_ext.h   |  2 ++
 tests/pgraph/test_loader_types.py               |  9 +++++++++
 4 files changed, 23 insertions(+)

diff --git a/panda/src/pgraph/loaderFileTypeRegistry.h b/panda/src/pgraph/loaderFileTypeRegistry.h
index cccb3af333..27f29664dd 100644
--- a/panda/src/pgraph/loaderFileTypeRegistry.h
+++ b/panda/src/pgraph/loaderFileTypeRegistry.h
@@ -53,6 +53,8 @@ PUBLISHED:
 
   static LoaderFileTypeRegistry *get_global_ptr();
 
+  EXTENSION(PyObject *__reduce__() const);
+
 private:
   void record_extension(const std::string &extension, LoaderFileType *type);
 
diff --git a/panda/src/pgraph/loaderFileTypeRegistry_ext.cxx b/panda/src/pgraph/loaderFileTypeRegistry_ext.cxx
index 20450e9195..eba66353df 100644
--- a/panda/src/pgraph/loaderFileTypeRegistry_ext.cxx
+++ b/panda/src/pgraph/loaderFileTypeRegistry_ext.cxx
@@ -18,6 +18,7 @@
 #include "pythonLoaderFileType.h"
 
 extern struct Dtool_PyTypedObject Dtool_LoaderFileType;
+extern struct Dtool_PyTypedObject Dtool_LoaderFileTypeRegistry;
 
 /**
  * Registers a loader file type that is implemented in Python.
@@ -112,4 +113,13 @@ unregister_type(PyObject *type) {
   Py_XDECREF(save_func);
 }
 
+/**
+ * Implements pickle support.
+ */
+PyObject *Extension<LoaderFileTypeRegistry>::
+__reduce__() const {
+  PyObject *func = PyObject_GetAttrString((PyObject *)&Dtool_LoaderFileTypeRegistry, "get_global_ptr");
+  return Py_BuildValue("N()", func);
+}
+
 #endif
diff --git a/panda/src/pgraph/loaderFileTypeRegistry_ext.h b/panda/src/pgraph/loaderFileTypeRegistry_ext.h
index 9c9815c20d..e5325101e6 100644
--- a/panda/src/pgraph/loaderFileTypeRegistry_ext.h
+++ b/panda/src/pgraph/loaderFileTypeRegistry_ext.h
@@ -33,6 +33,8 @@ public:
   void register_deferred_type(PyObject *entry_point);
 
   void unregister_type(PyObject *type);
+
+  PyObject *__reduce__() const;
 };
 
 #endif  // HAVE_PYTHON
diff --git a/tests/pgraph/test_loader_types.py b/tests/pgraph/test_loader_types.py
index 8fca1c5849..2a0114852f 100644
--- a/tests/pgraph/test_loader_types.py
+++ b/tests/pgraph/test_loader_types.py
@@ -3,6 +3,7 @@ import pytest
 import tempfile
 import os
 from contextlib import contextmanager
+import sys
 
 
 @pytest.fixture
@@ -218,3 +219,11 @@ def test_loader_ram_cache(test_filename):
         assert model1 == model2
 
         ModelPool.release_model(model2)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 4), reason="Requires Python 3.4")
+def test_loader_file_type_registry_pickle():
+    from direct.stdpy.pickle import dumps, loads
+
+    registry = LoaderFileTypeRegistry.get_global_ptr()
+    assert loads(dumps(registry, -1)) == registry