From 8d8cda5a03866601a9d65a547e22bbf2740b889e Mon Sep 17 00:00:00 2001
From: Gabriel Studer <gabriel.studer@unibas.ch>
Date: Thu, 1 Feb 2024 17:09:14 +0100
Subject: [PATCH] Convenience functions to get entity ids in MMCifInfo object
---
modules/io/pymod/export_mmcif_io.cc | 2 ++
modules/io/src/mol/mmcif_info.cc | 18 ++++++++++++++++++
modules/io/src/mol/mmcif_info.hh | 4 ++++
3 files changed, 24 insertions(+)
diff --git a/modules/io/pymod/export_mmcif_io.cc b/modules/io/pymod/export_mmcif_io.cc
index 325fb2439..3ca9d9dcd 100644
--- a/modules/io/pymod/export_mmcif_io.cc
+++ b/modules/io/pymod/export_mmcif_io.cc
@@ -550,6 +550,8 @@ void export_mmcif_io()
.def("GetEntityBranchChains", &MMCifInfo::GetEntityBranchChains)
.def("SetEntityDesc", &MMCifInfo::SetEntityDesc)
.def("GetEntityDesc", &MMCifInfo::GetEntityDesc, return_value_policy<copy_const_reference>())
+ .def("GetEntityIds", &MMCifInfo::GetEntityIds)
+ .def("GetEntityIdsOfType", &MMCifInfo::GetEntityIdsOfType)
.add_property("citations", make_function(&MMCifInfo::GetCitations,
return_value_policy<copy_const_reference>()))
.add_property("biounits", make_function(&MMCifInfo::GetBioUnits,
diff --git a/modules/io/src/mol/mmcif_info.cc b/modules/io/src/mol/mmcif_info.cc
index 96f256ca9..e723440a4 100644
--- a/modules/io/src/mol/mmcif_info.cc
+++ b/modules/io/src/mol/mmcif_info.cc
@@ -310,6 +310,24 @@ void MMCifInfo::SetEntityDesc(const String& entity_id,
entity_desc_[entity_id] = entity_desc;
}
+std::vector<String> MMCifInfo::GetEntityIds() const {
+ std::vector<String> vec;
+ for(auto it: entity_desc_) {
+ vec.push_back(it.first);
+ }
+ return vec;
+}
+
+std::vector<String> MMCifInfo::GetEntityIdsOfType(const String& entity_type) const {
+ std::vector<String> vec;
+ for(auto it: entity_desc_) {
+ if(it.second.entity_type == entity_type) {
+ vec.push_back(it.first);
+ }
+ }
+ return vec;
+}
+
std::ostream& operator<<(std::ostream& os, const MMCifInfoEntityBranchLink& eb)
{
os << "<MMCifInfoEntityBranchLink atom1:" << eb.GetAtom1() << " atom2:"
diff --git a/modules/io/src/mol/mmcif_info.hh b/modules/io/src/mol/mmcif_info.hh
index d4ca56194..438b90611 100644
--- a/modules/io/src/mol/mmcif_info.hh
+++ b/modules/io/src/mol/mmcif_info.hh
@@ -1233,6 +1233,10 @@ public:
void SetEntityDesc(const String& entity_id,
const MMCifEntityDesc& entity_desc);
+ std::vector<String> GetEntityIds() const;
+
+ std::vector<String> GetEntityIdsOfType(const String& type) const;
+
//protected:
private:
--
GitLab