Change to set_structure. Allows writing of individual (S)MCRAs

Bio/PDB/PDBIO.py

@@ -5,6 +5,7 @@
 
 """Output of PDB files."""
 
+from Bio.PDB.StructureBuilder import StructureBuilder # To allow saving of chains, residues, etc..
 from Bio.Data.IUPACData import atom_weights # Allowed Elements
 
 _ATOM_FORMAT_STRING="%s%5i %-4s%c%3s %c%4i%c   %8.3f%8.3f%8.3f%6.2f%6.2f      %4s%2s%2s\n"
@@ -91,10 +92,46 @@ def _get_atom_line(self, atom, hetfield, segid, atom_number, resname,
 
     # Public methods
 
-    def set_structure(self, structure):
+    def set_structure(self, pdb_object):
+        # Check what the user is providing and build a structure appropriately
+        if pdb_object.level == "S":
+            structure = pdb_object
+        else:
+            sb = StructureBuilder()
+            sb.init_structure('pdb')
+            sb.init_seg(' ')
+            # Build parts as necessary
+            if pdb_object.level == "M":
+                sb.structure.add(pdb_object)
+                self.structure = sb.structure
+            else:
+                sb.init_model(0)
+                if pdb_object.level == "C":
+                    sb.structure[0].add(pdb_object)
+                else:
+                    sb.init_chain('A')
+                    if pdb_object.level == "R":
+                        try:
+                            parent_id = pdb_object.parent.id
+                            sb.structure[0]['A'].id = parent_id
+                        except Exception:
+                            pass
+                        sb.structure[0]['A'].add(pdb_object)
+                    else:
+                        # Atom
+                        sb.init_residue('DUM', ' ', 1, ' ')
+                        try:
+                            parent_id = pdb_object.parent.parent.id
+                            sb.structure[0]['A'].id = parent_id
+                        except Exception:
+                            pass
+                        sb.structure[0]['A'].child_list[0].add(pdb_object)
+
+            # Return structure
+            structure = sb.structure
         self.structure=structure
 
-    def save(self, file, select=Select(), write_end=0):
+    def save(self, file, select=Select(), write_end=1):
         """
         @param file: output file
         @type file: string or filehandle