Fix Python container and pickle API

openPMD · Jun 6, 2024 · 4ae7a3f · 4ae7a3f
1 parent 8c57285
commit 4ae7a3f
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/include/openPMD/binding/python/Container.H b/include/openPMD/binding/python/Container.H
@@ -34,6 +34,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <pybind11/attr.h>
 #include <sstream>
 #include <string>
 #include <utility>
@@ -118,11 +119,12 @@ Class_ finalize_container(Class_ cl)
     // keep same policy as Container class: missing keys are created
     cl.def(
         "__getitem__",
-        [](Map &m, KeyType const &k) -> MappedType & { return m[k]; },
+        [](Map &m, KeyType const &k) -> MappedType { return m[k]; },
         // copy + keepalive
         // All objects in the openPMD object model are handles, so using a copy
         // is safer and still performant.
-        py::return_value_policy::copy);
+        py::return_value_policy::move,
+        py::keep_alive<0, 1>());
 
     // Assignment provided only if the type is copyable
     py::detail::map_assignment<Map, Class_>(cl);

diff --git a/include/openPMD/binding/python/Pickle.hpp b/include/openPMD/binding/python/Pickle.hpp
@@ -58,7 +58,7 @@ add_pickle(pybind11::class_<T_Args...> &cl, T_SeriesAccessor &&seriesAccessor)
         },
 
         // __setstate__
-        [&seriesAccessor](py::tuple t) {
+        [&seriesAccessor](py::tuple const &t) {
             // our tuple has exactly two elements: filePath & group
             if (t.size() != 2)
                 throw std::runtime_error("Invalid state!");

diff --git a/src/binding/python/MeshRecordComponent.cpp b/src/binding/python/MeshRecordComponent.cpp
@@ -84,7 +84,10 @@ void init_MeshRecordComponent(py::module &m)
     add_pickle(
         cl, [](openPMD::Series &series, std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
-            return series.iterations[n_it].meshes[group.at(3)][group.at(4)];
+            return series.iterations[n_it]
+                .meshes[group.at(3)]
+                       [group.size() < 5 ? MeshRecordComponent::SCALAR
+                                         : group.at(4)];
         });
 
     finalize_container<PyMeshRecordComponentContainer>(py_mrc_cnt);

diff --git a/src/binding/python/RecordComponent.cpp b/src/binding/python/RecordComponent.cpp
@@ -1124,8 +1124,8 @@ void init_RecordComponent(py::module &m)
     add_pickle(
         cl, [](openPMD::Series &series, std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));
-            return series.iterations[n_it]
-                .particles[group.at(3)][group.at(4)][group.at(5)];
+            return series.iterations[n_it].particles[group.at(3)][group.at(
+                4)][group.size() < 6 ? RecordComponent::SCALAR : group.at(5)];
         });
 
     addRecordComponentSetGet(cl);