Add support for pickling string/array Fields and Keys

lsst · Jun 4, 2019 · a4aef40 · a4aef40
1 parent edeebbf
commit a4aef40
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 9 deletions.
diff --git a/include/lsst/afw/table/detail/Access.h b/include/lsst/afw/table/detail/Access.h
@@ -75,6 +75,17 @@ class Access final {
     /// @internal Access to the private Key constructor.
     static Key<Flag> makeKey(int offset, int bit) { return Key<Flag>(offset, bit); }
 
+    /// @internal Access to the private Key constructor.
+    static Key<std::string> makeKeyString(int offset, int size) {
+        return Key<std::string>(offset, FieldBase<std::string>(size));
+    }
+
+    /// @internal Access to the private Key constructor.
+    template <typename U>
+    static Key<Array<U>> makeKeyArray(int offset, int size) {
+        return Key<Array<U>>(offset, FieldBase<Array<U>>(size));
+    }
+
     /// @internal Add some padding to a schema without adding a field.
     static void padSchema(Schema &schema, int bytes) {
         schema._edit();

diff --git a/python/lsst/afw/table/schema/schema.cc b/python/lsst/afw/table/schema/schema.cc
@@ -82,6 +82,48 @@ void declareFieldBaseSpecializations(PyFieldBase<std::string> &cls) {
     cls.def("getSize", &FieldBase<std::string>::getSize);
 }
 
+// Specializations for Field
+
+template <typename T>
+void declareFieldSpecializations(PyField<T> &cls) {
+    cls.def(py::pickle(
+            [](Field<T> const &self) {
+                /* Return a tuple that fully encodes the state of the object */
+                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits());
+            },
+            [](py::tuple t) {
+                if (t.size() != 3) throw std::runtime_error("Invalid number of parameters when unpickling!");
+                return Field<T>(t[0].cast<std::string>(), t[1].cast<std::string>(), t[2].cast<std::string>());
+            }));
+}
+
+template <typename U>
+void declareFieldSpecializations(PyField<Array<U>> &cls) {
+    cls.def(py::pickle(
+            [](Field<Array<U>> const &self) {
+                /* Return a tuple that fully encodes the state of the object */
+                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits(), self.getSize());
+            },
+            [](py::tuple t) {
+                if (t.size() != 4) throw std::runtime_error("Invalid number of parameters when unpickling!");
+                return Field<Array<U>>(t[0].cast<std::string>(), t[1].cast<std::string>(),
+                                       t[2].cast<std::string>(), t[3].cast<int>());
+            }));
+}
+
+void declareFieldSpecializations(PyField<std::string> &cls) {
+    cls.def(py::pickle(
+            [](Field<std::string> const &self) {
+                /* Return a tuple that fully encodes the state of the object */
+                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits(), self.getSize());
+            },
+            [](py::tuple t) {
+                if (t.size() != 4) throw std::runtime_error("Invalid number of parameters when unpickling!");
+                return Field<std::string>(t[0].cast<std::string>(), t[1].cast<std::string>(),
+                                          t[2].cast<std::string>(), t[3].cast<int>());
+            }));
+}
+
 // Specializations for KeyBase
 
 template <typename T>
@@ -195,6 +237,30 @@ void declareKeySpecializations(PyKey<Array<U>> &cls) {
         }
         return py::tuple(result);
     });
+    cls.def(py::pickle(
+            [](Key<Array<U>> const &self) {
+                /* Return a tuple that fully encodes the state of the object */
+                return py::make_tuple(self.getOffset(), self.getElementCount());
+            },
+            [](py::tuple t) {
+                if (t.size() != 2) throw std::runtime_error("Invalid number of parameters when unpickling!");
+                return detail::Access::makeKeyArray<U>(t[0].cast<int>(), t[1].cast<int>());
+            }));
+}
+
+void declareKeySpecializations(PyKey<std::string> &cls) {
+    declareKeyAccessors(cls);
+    cls.def_property_readonly("subfields", [](py::object const &) { return py::none(); });
+    cls.def_property_readonly("subkeys", [](py::object const &) { return py::none(); });
+    cls.def(py::pickle(
+            [](Key<std::string> const &self) {
+                /* Return a tuple that fully encodes the state of the object */
+                return py::make_tuple(self.getOffset(), self.getElementCount());
+            },
+            [](py::tuple t) {
+                if (t.size() != 2) throw std::runtime_error("Invalid number of parameters when unpickling!");
+                return detail::Access::makeKeyString(t[0].cast<int>(), t[1].cast<int>());
+            }));
 }
 
 // Wrap all helper classes (FieldBase, KeyBase, Key, Field, SchemaItem) declarefor a Schema field type.
@@ -217,6 +283,8 @@ void declareSchemaType(py::module &mod) {
 
     // Field
     PyField<T> clsField(mod, ("Field" + suffix).c_str());
+    declareFieldSpecializations(clsField);
+
     mod.attr("_Field")[pySuffix] = clsField;
     clsField.def(py::init([astropyUnit](  // capture by value to refcount in Python instead of dangle in C++
                                   std::string const &name, std::string const &doc, py::str const &units,
@@ -240,15 +308,6 @@ void declareSchemaType(py::module &mod) {
     clsField.def("copyRenamed", &Field<T>::copyRenamed);
     utils::python::addOutputOp(clsField, "__str__");
     utils::python::addOutputOp(clsField, "__repr__");
-    clsField.def(py::pickle(
-            [](Field<T> const &self) {
-                /* Return a tuple that fully encodes the state of the object */
-                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits());
-            },
-            [](py::tuple t) {
-                if (t.size() != 3) throw std::runtime_error("Invalid number of parameters when unpickling!");
-                return Field<T>(t[0].cast<std::string>(), t[1].cast<std::string>(), t[2].cast<std::string>());
-            }));
 
     // Key
     PyKey<T> clsKey(mod, ("Key" + suffix).c_str());

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -184,6 +184,11 @@ def testPickle(self):
         schema.addField("i", type="I", doc="int")
         schema.addField("f", type="F", doc="float", units="m2")
         schema.addField("flag", type="Flag", doc="a flag")
+        schema.addField("string", type="String", doc="A string field", size=42)
+        schema.addField("variable_string", type="String", doc="A variable-length string field", size=0)
+        schema.addField("array", type="ArrayF", doc="An array field", size=10)
+        schema.addField("variable_array", type="ArrayF", doc="A variable-length array field", size=0)
+
         pickled = pickle.dumps(schema, protocol=pickle.HIGHEST_PROTOCOL)
         unpickled = pickle.loads(pickled)
         self.assertEqual(schema, unpickled)
@@ -410,6 +415,10 @@ def testPickle(self):
         schema.addField("i", type="I", doc="int")
         schema.addField("f", type="F", doc="float", units="m2")
         schema.addField("flag", type="Flag", doc="a flag")
+        schema.addField("string", type="String", doc="A string field", size=42)
+        schema.addField("variable_string", type="String", doc="A variable-length string field", size=0)
+        schema.addField("array", type="ArrayF", doc="An array field", size=10)
+        schema.addField("variable_array", type="ArrayF", doc="A variable-length array field", size=0)
         mapper = lsst.afw.table.SchemaMapper(schema)
         mapper.addMinimalSchema(schema)
         inKey = schema.addField("bb", type=float)