From 5cba0319fe43ea322f9f7668cd460364399c43b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Franz=20P=C3=B6schel?= <franz.poeschel@gmail.com>
Date: Mon, 17 Jun 2024 12:05:30 +0200
Subject: [PATCH] Add pickle support for Series and Iteration

---
 include/openPMD/Iteration.hpp       | 7 +++++++
 src/backend/Attributable.cpp        | 1 +
 test/python/unittest/API/APITest.py | 8 ++++++++
 3 files changed, 16 insertions(+)

diff --git a/include/openPMD/Iteration.hpp b/include/openPMD/Iteration.hpp
index a8f4d7e43d..52bf43293a 100644
--- a/include/openPMD/Iteration.hpp
+++ b/include/openPMD/Iteration.hpp
@@ -131,6 +131,8 @@ class Iteration : public Attributable
     friend class WriteIterations;
     friend class SeriesIterator;
     friend class internal::AttributableData;
+    template <typename T>
+    friend T &internal::makeOwning(T &self, Series);
 
 public:
     Iteration(Iteration const &) = default;
@@ -258,6 +260,11 @@ class Iteration : public Attributable
         return *m_iterationData;
     }
 
+    inline std::shared_ptr<Data_t> getShared()
+    {
+        return m_iterationData;
+    }
+
     inline void setData(std::shared_ptr<Data_t> data)
     {
         m_iterationData = std::move(data);
diff --git a/src/backend/Attributable.cpp b/src/backend/Attributable.cpp
index 2a75844d75..a4d993b9f9 100644
--- a/src/backend/Attributable.cpp
+++ b/src/backend/Attributable.cpp
@@ -528,5 +528,6 @@ namespace internal
     template Mesh &makeOwning(Mesh &, Series);
     template Record &makeOwning(Record &, Series);
     template ParticleSpecies &makeOwning(ParticleSpecies &, Series);
+    template Iteration &makeOwning(Iteration &, Series);
 } // namespace internal
 } // namespace openPMD
diff --git a/test/python/unittest/API/APITest.py b/test/python/unittest/API/APITest.py
index 6ff987f657..59e6b5c97e 100644
--- a/test/python/unittest/API/APITest.py
+++ b/test/python/unittest/API/APITest.py
@@ -971,6 +971,8 @@ def testPickle(self):
         series.flush()
 
         # Pickle
+        pickled_s = pickle.dumps(series)
+        pickled_i = pickle.dumps(i)
         pickled_E = pickle.dumps(E)
         pickled_E_x = pickle.dumps(E_x)
         pickled_electrons = pickle.dumps(electrons)
@@ -980,6 +982,7 @@ def testPickle(self):
         pickled_w = pickle.dumps(w)
         print(f"This is my pickled object:\n{pickled_E_x}\n")
 
+        series.close()
         del E
         del E_x
         del electrons
@@ -987,9 +990,12 @@ def testPickle(self):
         del pos
         del pos_y
         del w
+        del i
         del series
 
         # Unpickling the object
+        series = pickle.loads(pickled_s)
+        i = pickle.loads(pickled_i)
         E = pickle.loads(pickled_E)
         E_x = pickle.loads(pickled_E_x)
         electrons = pickle.loads(pickled_electrons)
@@ -1000,6 +1006,8 @@ def testPickle(self):
         print(
             f"This is E_x.position of the unpickled object:\n{E_x.position}\n")
 
+        self.assertIsInstance(series, io.Series)
+        self.assertIsInstance(i, io.Iteration)
         self.assertIsInstance(E, io.Mesh)
         self.assertIsInstance(E_x, io.Mesh_Record_Component)
         self.assertIsInstance(electrons, io.ParticleSpecies)