serializing (pickling in Python) of cif.Document and cif.Block (#258)

Author: wojdyr

include/gemmi/cifdoc.hpp

@@ -522,6 +522,7 @@ struct Item {
     Block frame;
   };
 
+  Item() : type(ItemType::Erased) {}
   explicit Item(LoopArg)
     : type{ItemType::Loop}, loop{} {}
   explicit Item(std::string&& t)

include/gemmi/serialize.hpp

@@ -8,6 +8,7 @@
 #define GEMMI_SERIALIZE_HPP_
 
 #include "model.hpp"
+#include "cifdoc.hpp"
 
 #define SERIALIZE(Struct, ...) \
 template <typename Archive> \
@@ -144,6 +145,38 @@ SERIALIZE(Structure, o.name, o.cell, o.spacegroup_hm, o.models,
           o.input_format, o.has_d_fraction, o.ter_status,
           o.has_origx, o.origx, o.info, o.shortened_ccd_codes,
           o.raw_remarks, o.resolution)
+
+
+namespace cif {
+
+SERIALIZE(Loop, o.tags, o.values)
+SERIALIZE(Block, o.name, o.items)
+SERIALIZE(Document, o.source, o.blocks)
+
+template <typename Archive>
+void serialize(Archive& archive, Item& o) {
+  archive(o.type, o.line_number);
+  switch (o.type) {
+    case ItemType::Pair:
+    case ItemType::Comment: new(&o.pair) cif::Pair; archive(o.pair); break;
+    case ItemType::Loop: new(&o.loop) cif::Loop; archive(o.loop); break;
+    case ItemType::Frame: new(&o.frame) cif::Block; archive(o.frame); break;
+    case ItemType::Erased: break;
+  }
+}
+template <typename Archive>
+void serialize(Archive& archive, const Item& o) {
+  archive(o.type, o.line_number);
+  switch (o.type) {
+    case ItemType::Pair:
+    case ItemType::Comment: archive(o.pair); break;
+    case ItemType::Loop: archive(o.loop); break;
+    case ItemType::Frame: archive(o.frame); break;
+    case ItemType::Erased: break;
+  }
+}
+
+} // namespace cif
 } // namespace gemmi
 
 #undef SERIALIZE

python/cif.cpp

@@ -9,6 +9,7 @@
 #include "gemmi/read_cif.hpp"  // for read_cif_gz
 
 #include "common.h"
+#include "serial.h"  // for getstate, setstate
 #include "make_iterator.h"
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
@@ -168,6 +169,8 @@ void add_cif(nb::module_& cif) {
         return os.str();
     }, nb::arg("mmjson")=false, nb::arg("lowercase_names")=true,
     "Returns JSON representation in a string.")
+    .def("__getstate__", &getstate<Document>)
+    .def("__setstate__", &setstate<Document>)
     .def("__repr__", [](const Document &d) {
         std::string s = "<gemmi.cif.Document with ";
         s += std::to_string(d.blocks.size());
@@ -297,6 +300,8 @@ void add_cif(nb::module_& cif) {
         write_cif_block_to_stream(os, self, opt);
         return os.str();
     }, nb::arg("options")=WriteOptions(), "Returns a string in CIF format.")
+    .def("__getstate__", &getstate<Block>)
+    .def("__setstate__", &setstate<Block>)
     .def("__repr__", [](const Block &self) {
         return gemmi::cat("<gemmi.cif.Block ", self.name, '>');
     });

tests/test_cif.py

@@ -2,11 +2,12 @@
 
 import gc
 import os
+import pickle
 import unittest
 from gemmi import cif
 
 class TestDoc(unittest.TestCase):
-    def test_slice(self):
+    def test_slicing_and_pickling(self):
         doc = cif.read_string("""
             data_a
             _one 1 _two 2 _three 3
@@ -15,23 +16,16 @@ def test_slice(self):
             data_c
             _two 2 _four 4 _six 6
         """)
+        self.assertTrue('a' in doc)
+        self.assertFalse('d' in doc)
         self.assertEqual([b.name for b in doc[:1]], ['a'])
         self.assertEqual([b.name for b in doc[1:]], ['b', 'c'])
         self.assertEqual([b.name for b in doc[:]], ['a', 'b', 'c'])
         self.assertEqual([b.name for b in doc[1:-1]], ['b'])
         self.assertEqual([b.name for b in doc[1:1]], [])
 
-    def test_contains(self):
-        doc = cif.read_string("""
-            data_a
-            _one 1 _two 2 _three 3
-            data_b
-            _four 4
-            data_c
-            _two 2 _four 4 _six 6
-        """)
-        self.assertEqual('a' in doc, True)
-        self.assertEqual('d' in doc, False)
+        unpickled = pickle.loads(pickle.dumps(doc))
+        self.assertEqual(unpickled.as_string(), doc.as_string())
 
 class TestBlock(unittest.TestCase):
     def test_find(self):
@@ -358,7 +352,11 @@ def test_case_sensitivity(self):
             _One 1 _thrEE 3
             _NonLoop_a alpha
             loop_ _lbBb _ln  B 1  D 2"""
-        self.assertEqual(block.as_string().split(), expected.split())
+        block_str = block.as_string()
+        self.assertEqual(block_str.split(), expected.split())
+
+        unpickled = pickle.loads(pickle.dumps(block))
+        self.assertEqual(unpickled.as_string(), block_str)
 
 class TestQuote(unittest.TestCase):
     def test_quote(self):