Add pickle support for Series and Iteration

franzpoeschel · Jul 19, 2024 · 5cba031 · 5cba031
1 parent 71c7f75
commit 5cba031
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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,16 +982,20 @@ 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
         del momentum
         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)