diff --git a/python/include/nuri/python/utils.h b/python/include/nuri/python/utils.h index a9370294..3936895a 100644 --- a/python/include/nuri/python/utils.h +++ b/python/include/nuri/python/utils.h @@ -493,7 +493,9 @@ py::class_ &add_sequence_interface(py::class_ &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), kReturnsSubobject, py::arg("idx")); diff --git a/python/src/nuri/fmt/cif.cpp b/python/src/nuri/fmt/cif.cpp index d141464b..31f802f5 100644 --- 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::str(*val).cast()); + const internal::CifValue &val = table[i][j]; + data[i].push_back(val.is_null() ? py::none().cast() + : py::str(*val).cast()); } } } @@ -289,35 +294,40 @@ cif_ddl2_frame_as_dict(const PyCifFrame &frame) { void bind_cif(py::module &m) { PyCifTableIterator::bind(m); - py::class_(m, "CifTable") // + py::class_(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_(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_ 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( m, "_CifFrameList", "_CifFrameList index out of range"); - py::class_(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_ 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_(m, "CifParser") .def("__iter__", pass_through) diff --git a/python/test/fmt/cif_test.py b/python/test/fmt/cif_test.py index f2b9f14d..d9439123 100644 --- 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"