Merge pull request #447 from seoklab/python/fmt/fix-cif

fix(python/fmt/cif): pass internal reference to subobject and support inconsistent CIF files
seoklab · Jan 6, 2025 · a43aaed · a43aaed
2 parents 81f797b + b7e504e
commit a43aaed
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 28 deletions.
diff --git a/python/include/nuri/python/utils.h b/python/include/nuri/python/utils.h
@@ -493,7 +493,9 @@ py::class_<T> &add_sequence_interface(py::class_<T> &cls, Size size,
   cls.def("__len__", size);
   cls.def(
       "__contains__",
-      [size](T &self, int idx) { return 0 <= idx && idx < size(self); },
+      [size](T &self, int idx) {
+        return 0 <= idx && idx < std::invoke(size, self);
+      },
       py::arg("idx"));
   cls.def("__getitem__", std::forward<Getter>(getter), kReturnsSubobject,
           py::arg("idx"));

diff --git a/python/src/nuri/fmt/cif.cpp b/python/src/nuri/fmt/cif.cpp
@@ -249,22 +249,27 @@ cif_ddl2_frame_as_dict(const PyCifFrame &frame) {
 
     slots.clear();
     slots.reserve(table.cols());
-    for (std::string_view pk: parent_keys)
-      slots.push_back(grouped.find(pk));
+    for (std::string_view pk: parent_keys) {
+      auto it = grouped.find(pk);
+      ABSL_DCHECK(it != grouped.end());
+      slots.push_back(it);
 
-    for (const auto &row: table) {
-      for (int i = 0; i < table.cols(); ++i) {
-        auto it = slots[i];
-        ABSL_DCHECK(it != grouped.end());
+      auto &data = it->second.second;
+      data.resize(table.size());
 
-        auto &data = it->second.second;
+      for (auto &row: data)
+        row.reserve(row.size() + it->second.first.size());
+    }
+
+    for (int i = 0; i < table.size(); ++i) {
+      for (int j = 0; j < table.cols(); ++j) {
+        auto it = slots[j];
 
-        if (data.empty() || data.back().size() == it->second.first.size())
-          data.emplace_back().reserve(it->second.first.size());
+        auto &data = it->second.second;
 
-        const internal::CifValue &val = row[i];
-        data.back().push_back(val.is_null() ? py::none().cast<py::object>()
-                                            : py::str(*val).cast<py::object>());
+        const internal::CifValue &val = table[i][j];
+        data[i].push_back(val.is_null() ? py::none().cast<py::object>()
+                                        : py::str(*val).cast<py::object>());
       }
     }
   }
@@ -289,35 +294,40 @@ cif_ddl2_frame_as_dict(const PyCifFrame &frame) {
 void bind_cif(py::module &m) {
   PyCifTableIterator::bind(m);
 
-  py::class_<PyCifTable>(m, "CifTable")  //
+  py::class_<PyCifTable>(m, "CifTable")
       .def("__iter__", &PyCifTable::iter, kReturnsSubobject)
-      .def("__getitem__", &PyCifTable::get)
+      .def("__getitem__", &PyCifTable::get, py::arg("idx"))
       .def("__len__", &PyCifTable::size)
+      .def(
+          "__contains__",
+          [](const PyCifTable &self, int idx) {
+            return 0 <= idx && idx < self.size();
+          },
+          py::arg("idx"))
       .def("keys", &PyCifTable::keys);
 
   PyCifFrameIterator::bind(m);
 
-  py::class_<PyCifFrame>(m, "CifFrame")  //
-      .def("__iter__", &PyCifFrame::iter, kReturnsSubobject)
-      .def("__getitem__", &PyCifFrame::get, kReturnsSubobject)
-      .def("__len__", &PyCifFrame::size)
-      .def("prefix_search_first", &PyCifFrame::prefix_search_first,
-           py::arg("prefix"), kReturnsSubobject, R"doc(
+  py::class_<PyCifFrame> cf(m, "CifFrame");
+  add_sequence_interface(cf, &PyCifFrame::size, &PyCifFrame::get,
+                         &PyCifFrame::iter);
+  cf.def("prefix_search_first", &PyCifFrame::prefix_search_first,
+         py::arg("prefix"), kReturnsSubobject, R"doc(
 Search for the first table containing a column starting with the given prefix.
 
 :param prefix: The prefix to search for.
 :return: The first table containing the given prefix, or None if not found.
-)doc")
-      .def_property_readonly("name", &PyCifFrame::name);
+)doc");
+  cf.def_property_readonly("name", &PyCifFrame::name);
 
   bind_opaque_vector<internal::CifFrame, PyCifFrame, wrap_cif_frame>(
       m, "_CifFrameList", "_CifFrameList index out of range");
 
-  py::class_<PyCifBlock>(m, "CifBlock")  //
-      .def_property_readonly("data", &PyCifBlock::data)
-      .def_property_readonly("name", &PyCifBlock::name)
-      .def_property_readonly("save_frames", &PyCifBlock::save_frames)
-      .def_property_readonly("is_global", &PyCifBlock::is_global);
+  py::class_<PyCifBlock> cb(m, "CifBlock");
+  cb.def_property_readonly("name", &PyCifBlock::name)
+      .def_property_readonly("is_global", &PyCifBlock::is_global)
+      .def_property_readonly("save_frames", &PyCifBlock::save_frames);
+  def_property_readonly_subobject(cb, "data", &PyCifBlock::data);
 
   py::class_<PyCifParser>(m, "CifParser")
       .def("__iter__", pass_through<PyCifParser>)

diff --git a/python/test/fmt/cif_test.py b/python/test/fmt/cif_test.py
@@ -81,6 +81,11 @@ def test_read_cif(test_data: Path):
     assert nonexistent is None
 
 
+def test_read_cif_temporary(test_data: Path):
+    data = next(read_cif(test_data / "1a8o.cif")).data
+    assert data.name == "1A8O"
+
+
 def test_convert_ddl2_cif(test_data: Path):
     cif = test_data / "1a8o.cif"