1 files changed, 72 insertions, 25 deletions

diff --git a/ext/pybind11/tests/test_pickling.cpp b/ext/pybind11/tests/test_pickling.cpp
index 52b1dbc30..9dc63bda3 100644
--- a/ext/pybind11/tests/test_pickling.cpp
+++ b/ext/pybind11/tests/test_pickling.cpp
@@ -9,30 +9,28 @@
 
 #include "pybind11_tests.h"
 
-class Pickleable {
-public:
-    Pickleable(const std::string &value) : m_value(value) { }
-    const std::string &value() const { return m_value; }
-
-    void setExtra1(int extra1) { m_extra1 = extra1; }
-    void setExtra2(int extra2) { m_extra2 = extra2; }
-    int extra1() const { return m_extra1; }
-    int extra2() const { return m_extra2; }
-private:
-    std::string m_value;
-    int m_extra1 = 0;
-    int m_extra2 = 0;
-};
-
-class PickleableWithDict {
-public:
-    PickleableWithDict(const std::string &value) : value(value) { }
-
-    std::string value;
-    int extra;
-};
-
-test_initializer pickling([](py::module &m) {
+TEST_SUBMODULE(pickling, m) {
+    // test_roundtrip
+    class Pickleable {
+    public:
+        Pickleable(const std::string &value) : m_value(value) { }
+        const std::string &value() const { return m_value; }
+
+        void setExtra1(int extra1) { m_extra1 = extra1; }
+        void setExtra2(int extra2) { m_extra2 = extra2; }
+        int extra1() const { return m_extra1; }
+        int extra2() const { return m_extra2; }
+    private:
+        std::string m_value;
+        int m_extra1 = 0;
+        int m_extra2 = 0;
+    };
+
+    class PickleableNew : public Pickleable {
+    public:
+        using Pickleable::Pickleable;
+    };
+
     py::class_<Pickleable>(m, "Pickleable")
         .def(py::init<std::string>())
         .def("value", &Pickleable::value)
@@ -57,7 +55,38 @@ test_initializer pickling([](py::module &m) {
             p.setExtra2(t[2].cast<int>());
         });
 
+    py::class_<PickleableNew, Pickleable>(m, "PickleableNew")
+        .def(py::init<std::string>())
+        .def(py::pickle(
+            [](const PickleableNew &p) {
+                return py::make_tuple(p.value(), p.extra1(), p.extra2());
+            },
+            [](py::tuple t) {
+                if (t.size() != 3)
+                    throw std::runtime_error("Invalid state!");
+                auto p = PickleableNew(t[0].cast<std::string>());
+
+                p.setExtra1(t[1].cast<int>());
+                p.setExtra2(t[2].cast<int>());
+                return p;
+            }
+        ));
+
 #if !defined(PYPY_VERSION)
+    // test_roundtrip_with_dict
+    class PickleableWithDict {
+    public:
+        PickleableWithDict(const std::string &value) : value(value) { }
+
+        std::string value;
+        int extra;
+    };
+
+    class PickleableWithDictNew : public PickleableWithDict {
+    public:
+        using PickleableWithDict::PickleableWithDict;
+    };
+
     py::class_<PickleableWithDict>(m, "PickleableWithDict", py::dynamic_attr())
         .def(py::init<std::string>())
         .def_readwrite("value", &PickleableWithDict::value)
@@ -79,5 +108,23 @@ test_initializer pickling([](py::module &m) {
             /* Assign Python state */
             self.attr("__dict__") = t[2];
         });
+
+    py::class_<PickleableWithDictNew, PickleableWithDict>(m, "PickleableWithDictNew")
+        .def(py::init<std::string>())
+        .def(py::pickle(
+            [](py::object self) {
+                return py::make_tuple(self.attr("value"), self.attr("extra"), self.attr("__dict__"));
+            },
+            [](const py::tuple &t) {
+                if (t.size() != 3)
+                    throw std::runtime_error("Invalid state!");
+
+                auto cpp_state = PickleableWithDictNew(t[0].cast<std::string>());
+                cpp_state.extra = t[1].cast<int>();
+
+                auto py_state = t[2].cast<py::dict>();
+                return std::make_pair(cpp_state, py_state);
+            }
+        ));
 #endif
-});
+}