SCHWED-4847: Export functions to handle branched entitites to Python

modules/io/pymod/export_mmcif_io.cc

@@ -340,6 +340,7 @@ void export_mmcif_io()
     .def("ConnectBranchLink", &MMCifInfoEntityBranch::ConnectBranchLink)
     .def("SetAtom1", &MMCifInfoEntityBranch::SetAtom1)
     .def("SetAtom2", &MMCifInfoEntityBranch::SetAtom2)
+    .def(self_ns::str(self))
     .add_property("atom1", &MMCifInfoEntityBranch::GetAtom1,
                   &MMCifInfoEntityBranch::SetAtom1)
     .add_property("atom2", &MMCifInfoEntityBranch::GetAtom2,
@@ -350,6 +351,12 @@ void export_mmcif_io()
     .def(map_indexing_suite<MMCifInfoEntityBranchMap>())
   ;
 
+  class_<std::vector<MMCifInfoEntityBranch> >("MMCifInfoEntityBranchList",
+                                              init<>())
+    .def(vector_indexing_suite<std::vector<MMCifInfoEntityBranch> >())
+    .def(self_ns::str(self))
+  ;
+
   class_<MMCifInfo>("MMCifInfo", init<>())
     .def("AddCitation", &MMCifInfo::AddCitation)
     .def("GetCitations", make_function(&MMCifInfo::GetCitations,
@@ -384,6 +391,9 @@ void export_mmcif_io()
          (arg("num"), arg("date"), arg("status"), arg("major")=-1,
           arg("minor")=-1))
     .def("GetRevisions", &MMCifInfo::GetRevisions)
+    .def("AddEntityBranchLink", &MMCifInfo::AddEntityBranchLink)
+    .def("GetEntityBranchLinks", &MMCifInfo::GetEntityBranchLinks)
+    .def("ConnectBranchLinks", &MMCifInfo::ConnectBranchLinks)
     .add_property("citations", make_function(&MMCifInfo::GetCitations,
                                    return_value_policy<copy_const_reference>()))
     .add_property("biounits", make_function(&MMCifInfo::GetBioUnits,

modules/io/src/mol/mmcif_info.cc

@@ -239,4 +239,24 @@ void MMCifInfo::ConnectBranchLinks(mol::XCSEditor editor)
   }
 }
 
+std::ostream& operator<<(std::ostream& os, const MMCifInfoEntityBranch& eb)
+{
+  os << "<MMCifInfoEntityBranch atom1:" << eb.GetAtom1() << " atom2:"
+     << eb.GetAtom2() << ">";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os,
+                         const std::vector<MMCifInfoEntityBranch>& eb_list)
+{
+  os << "<MMCifInfoEntityBranchList";
+  std::vector<MMCifInfoEntityBranch>::const_iterator bl_it;
+  for (bl_it = eb_list.begin(); bl_it != eb_list.end(); ++bl_it) {
+    os << " <atom1:" << bl_it->GetAtom1() << " atom2:"
+       << bl_it->GetAtom2() << ">";
+  }
+  os << ">";
+  return os;
+}
+
 }} //ns

modules/io/src/mol/mmcif_info.hh

@@ -925,14 +925,27 @@ class DLLEXPORT_OST_IO MMCifInfoEntityBranch {
 public:
   MMCifInfoEntityBranch(mol::AtomHandle atom1, mol::AtomHandle atom2):
                 atom1_(atom1), atom2_(atom2) {}
-        mol::AtomHandle GetAtom1() const { return atom1_;}
-        mol::AtomHandle GetAtom2() const { return atom2_; }
-        void SetAtom1(mol::AtomHandle atom) { atom1_ = atom; }
-        void SetAtom2(mol::AtomHandle atom) { atom2_ = atom; }
-        void ConnectBranchLink(mol::XCSEditor editor)
-        {
-          editor.Connect(atom1_, atom2_);
-        }
+  mol::AtomHandle GetAtom1() const { return atom1_;}
+  mol::AtomHandle GetAtom2() const { return atom2_; }
+  void SetAtom1(mol::AtomHandle atom) { atom1_ = atom; }
+  void SetAtom2(mol::AtomHandle atom) { atom2_ = atom; }
+  void ConnectBranchLink(mol::XCSEditor editor) {
+    editor.Connect(atom1_, atom2_);
+  }
+
+  bool operator==(const MMCifInfoEntityBranch& eb) const {
+    if (this->atom1_ != eb.atom1_) {
+      return false;
+    }
+    if (this->atom2_ != eb.atom2_) {
+      return false;
+    }
+    return true;
+  }
+
+  bool operator!=(const MMCifInfoEntityBranch& eb) const {
+    return !this->operator == (eb);
+  }
 
 private:
   mol::AtomHandle atom1_;
@@ -1189,7 +1202,11 @@ private:
   std::map<String, std::vector<MMCifInfoEntityBranch> > entity_branches_;
 };
 
+DLLEXPORT_OST_IO std::ostream& operator<<(std::ostream& os, 
+                                          const MMCifInfoEntityBranch& eb);
 
+DLLEXPORT_OST_IO std::ostream& operator<<(std::ostream& os, 
+                             const std::vector<MMCifInfoEntityBranch>& eb_list);
 }} // ns
 
 #endif

modules/io/tests/test_io_mmcif.py

@@ -264,7 +264,7 @@ class TestMMCifInfo(unittest.TestCase):
     self.assertEqual(len(crambin_pdb.atoms), 327)
 
   def test_mmcifinfo_entitybranch(self):
-    # create dummy atoms
+    # test MMCifInfoEntityBranch
     eh = mol.CreateEntity()
     editor = eh.EditXCS();
     ch = editor.InsertChain("A");
@@ -278,6 +278,25 @@ class TestMMCifInfo(unittest.TestCase):
     branch.ConnectBranchLink(editor)
     self.assertEqual(atom2.GetBondPartners()[0].qualified_name, "A.MAN2.C1")
 
+    # test entity_branches_
+    ch = editor.InsertChain("B");
+    res1 = editor.AppendResidue(ch, "NAG");
+    res2 = editor.AppendResidue(ch, "NAG");
+    atom3 = editor.InsertAtom(res2, "C1", geom.Vec3());
+    atom4 = editor.InsertAtom(res1, "O4", geom.Vec3());
+    info = io.MMCifInfo()
+    info.AddEntityBranchLink("A", atom1, atom2)
+    info.AddEntityBranchLink(ch.name, atom3, atom4)
+
+    blinks = info.GetEntityBranchLinks()
+    self.assertEqual(blinks[0].GetAtom1().qualified_name, "A.MAN2.C1")
+    self.assertEqual(blinks[0].atom2.qualified_name, "A.BMA1.O3")
+    self.assertEqual(blinks[1].atom1.qualified_name, "B.NAG2.C1")
+    self.assertEqual(blinks[1].GetAtom2().qualified_name, "B.NAG1.O4")
+
+    info.ConnectBranchLinks(editor)
+    self.assertEqual(atom4.GetBondPartners()[0].qualified_name, "B.NAG2.C1")
+
 if __name__== '__main__':
   from ost import testutils
   testutils.RunTests()