1# Copyright 2009-2011 by Eric Talevich. All rights reserved. 2# Revisions copyright 2009-2013 by Peter Cock. All rights reserved. 3# Revisions copyright 2013 Lenna X. Peterson. All rights reserved. 4# Revisions copyright 2020 Joao Rodrigues. All rights reserved. 5# 6# Converted by Eric Talevich from an older unit test copyright 2002 7# by Thomas Hamelryck. 8# 9# This file is part of the Biopython distribution and governed by your 10# choice of the "Biopython License Agreement" or the "BSD 3-Clause License". 11# Please see the LICENSE file that should have been included as part of this 12# package. 13 14"""Unit tests for the Bio.PDB.PDBIO module.""" 15 16import os 17import tempfile 18import unittest 19import warnings 20 21from Bio import BiopythonWarning 22from Bio.PDB import PDBParser, PDBIO, Select 23from Bio.PDB import Atom, Residue 24from Bio.PDB.PDBExceptions import PDBConstructionException, PDBConstructionWarning 25 26 27class WriteTest(unittest.TestCase): 28 @classmethod 29 def setUpClass(self): 30 self.io = PDBIO() 31 self.parser = PDBParser(PERMISSIVE=1) 32 33 with warnings.catch_warnings(): 34 warnings.simplefilter("ignore", PDBConstructionWarning) 35 self.structure = self.parser.get_structure("example", "PDB/1A8O.pdb") 36 37 def test_pdbio_write_structure(self): 38 """Write a full structure using PDBIO.""" 39 struct1 = self.structure 40 # Ensure that set_structure doesn't alter parent 41 parent = struct1.parent 42 43 # Write full model to temp file 44 self.io.set_structure(struct1) 45 self.assertIs(parent, struct1.parent) 46 47 filenumber, filename = tempfile.mkstemp() 48 os.close(filenumber) 49 50 try: 51 self.io.save(filename) 52 53 struct2 = self.parser.get_structure("1a8o", filename) 54 nresidues = len(list(struct2.get_residues())) 55 56 self.assertEqual(len(struct2), 1) 57 self.assertEqual(nresidues, 158) 58 finally: 59 os.remove(filename) 60 61 def test_pdbio_write_preserve_numbering(self): 62 """Test writing PDB and preserve atom numbering.""" 63 self.io.set_structure(self.structure) 64 65 filenumber, filename = tempfile.mkstemp() 66 os.close(filenumber) 67 68 try: 69 self.io.save(filename) # default preserve_atom_numbering=False 70 71 struct = self.parser.get_structure("1a8o", filename) 72 serials = [a.serial_number for a in struct.get_atoms()] 73 og_serials = list(range(1, len(serials) + 1)) 74 self.assertEqual(og_serials, serials) 75 finally: 76 os.remove(filename) 77 78 def test_pdbio_write_auto_numbering(self): 79 """Test writing PDB and do not preserve atom numbering.""" 80 self.io.set_structure(self.structure) 81 82 filenumber, filename = tempfile.mkstemp() 83 os.close(filenumber) 84 85 try: 86 self.io.save(filename, preserve_atom_numbering=True) 87 88 struct = self.parser.get_structure("1a8o", filename) 89 serials = [a.serial_number for a in struct.get_atoms()] 90 og_serials = [a.serial_number for a in self.structure.get_atoms()] 91 92 self.assertEqual(og_serials, serials) 93 finally: 94 os.remove(filename) 95 96 def test_pdbio_write_residue(self): 97 """Write a single residue using PDBIO.""" 98 struct1 = self.structure 99 residue1 = list(struct1.get_residues())[0] 100 101 # Ensure that set_structure doesn't alter parent 102 parent = residue1.parent 103 104 # Write full model to temp file 105 self.io.set_structure(residue1) 106 self.assertIs(parent, residue1.parent) 107 filenumber, filename = tempfile.mkstemp() 108 os.close(filenumber) 109 try: 110 self.io.save(filename) 111 struct2 = self.parser.get_structure("1a8o", filename) 112 nresidues = len(list(struct2.get_residues())) 113 self.assertEqual(nresidues, 1) 114 finally: 115 os.remove(filename) 116 117 def test_pdbio_write_residue_w_chain(self): 118 """Write a single residue (chain id == X) using PDBIO.""" 119 struct1 = self.structure.copy() # make copy so we can change it 120 residue1 = list(struct1.get_residues())[0] 121 122 # Modify parent id 123 parent = residue1.parent 124 parent.id = "X" 125 126 # Write full model to temp file 127 self.io.set_structure(residue1) 128 filenumber, filename = tempfile.mkstemp() 129 os.close(filenumber) 130 try: 131 self.io.save(filename) 132 struct2 = self.parser.get_structure("1a8o", filename) 133 nresidues = len(list(struct2.get_residues())) 134 self.assertEqual(nresidues, 1) 135 136 # Assert chain remained the same 137 chain_id = [c.id for c in struct2.get_chains()][0] 138 self.assertEqual(chain_id, "X") 139 finally: 140 os.remove(filename) 141 142 def test_pdbio_write_residue_wout_chain(self): 143 """Write a single orphan residue using PDBIO.""" 144 struct1 = self.structure 145 residue1 = list(struct1.get_residues())[0] 146 147 residue1.parent = None # detach residue 148 149 # Write full model to temp file 150 self.io.set_structure(residue1) 151 152 filenumber, filename = tempfile.mkstemp() 153 os.close(filenumber) 154 try: 155 self.io.save(filename) 156 struct2 = self.parser.get_structure("1a8o", filename) 157 nresidues = len(list(struct2.get_residues())) 158 self.assertEqual(nresidues, 1) 159 160 # Assert chain is default: "A" 161 chain_id = [c.id for c in struct2.get_chains()][0] 162 self.assertEqual(chain_id, "A") 163 finally: 164 os.remove(filename) 165 166 def test_pdbio_write_custom_residue(self): 167 """Write a chainless residue using PDBIO.""" 168 res = Residue.Residue((" ", 1, " "), "DUM", "") 169 atm = Atom.Atom("CA", [0.1, 0.1, 0.1], 1.0, 1.0, " ", "CA", 1, "C") 170 res.add(atm) 171 172 # Ensure that set_structure doesn't alter parent 173 parent = res.parent 174 175 # Write full model to temp file 176 self.io.set_structure(res) 177 178 self.assertIs(parent, res.parent) 179 filenumber, filename = tempfile.mkstemp() 180 os.close(filenumber) 181 try: 182 self.io.save(filename) 183 struct2 = self.parser.get_structure("res", filename) 184 latoms = list(struct2.get_atoms()) 185 self.assertEqual(len(latoms), 1) 186 self.assertEqual(latoms[0].name, "CA") 187 self.assertEqual(latoms[0].parent.resname, "DUM") 188 self.assertEqual(latoms[0].parent.parent.id, "A") 189 finally: 190 os.remove(filename) 191 192 def test_pdbio_select(self): 193 """Write a selection of the structure using a Select subclass.""" 194 # Selection class to filter all alpha carbons 195 class CAonly(Select): 196 """Accepts only CA residues.""" 197 198 def accept_atom(self, atom): 199 if atom.name == "CA" and atom.element == "C": 200 return 1 201 202 struct1 = self.structure 203 # Ensure that set_structure doesn't alter parent 204 parent = struct1.parent 205 # Write to temp file 206 self.io.set_structure(struct1) 207 208 self.assertIs(parent, struct1.parent) 209 filenumber, filename = tempfile.mkstemp() 210 os.close(filenumber) 211 try: 212 self.io.save(filename, CAonly()) 213 struct2 = self.parser.get_structure("1a8o", filename) 214 nresidues = len(list(struct2.get_residues())) 215 self.assertEqual(nresidues, 70) 216 finally: 217 os.remove(filename) 218 219 def test_pdbio_missing_occupancy(self): 220 """Write PDB file with missing occupancy.""" 221 with warnings.catch_warnings(): 222 warnings.simplefilter("ignore", PDBConstructionWarning) 223 structure = self.parser.get_structure("test", "PDB/occupancy.pdb") 224 225 self.io.set_structure(structure) 226 filenumber, filename = tempfile.mkstemp() 227 os.close(filenumber) 228 try: 229 with warnings.catch_warnings(record=True) as w: 230 warnings.simplefilter("always", BiopythonWarning) 231 self.io.save(filename) 232 self.assertEqual(len(w), 1, w) 233 with warnings.catch_warnings(): 234 warnings.simplefilter("ignore", PDBConstructionWarning) 235 struct2 = self.parser.get_structure("test", filename) 236 atoms = struct2[0]["A"][(" ", 152, " ")] 237 self.assertIsNone(atoms["N"].get_occupancy()) 238 finally: 239 os.remove(filename) 240 241 def test_pdbio_write_truncated(self): 242 """Test parsing of truncated lines.""" 243 struct = self.structure 244 245 # Write to temp file 246 self.io.set_structure(struct) 247 filenumber, filename = tempfile.mkstemp() 248 os.close(filenumber) 249 try: 250 self.io.save(filename) 251 # Check if there are lines besides 'ATOM', 'TER' and 'END' 252 with open(filename) as handle: 253 record_set = {l[0:6] for l in handle} 254 record_set -= { 255 "ATOM ", 256 "HETATM", 257 "MODEL ", 258 "ENDMDL", 259 "TER\n", 260 "TER ", 261 "END\n", 262 "END ", 263 } 264 self.assertEqual(len(record_set), 0) 265 finally: 266 os.remove(filename) 267 268 def test_model_numbering(self): 269 """Preserve model serial numbers during I/O.""" 270 271 def confirm_numbering(struct): 272 self.assertEqual(len(struct), 3) 273 for idx, model in enumerate(struct): 274 self.assertEqual(model.serial_num, idx + 1) 275 self.assertEqual(model.serial_num, model.id + 1) 276 277 def confirm_single_end(fname): 278 """Ensure there is only one END statement in multi-model files.""" 279 with open(fname) as handle: 280 end_stment = [] 281 for iline, line in enumerate(handle): 282 if line.strip() == "END": 283 end_stment.append((line, iline)) 284 self.assertEqual(len(end_stment), 1) # Only one? 285 self.assertEqual(end_stment[0][1], iline) # Last line of the file? 286 287 with warnings.catch_warnings(): 288 warnings.simplefilter("ignore", PDBConstructionWarning) 289 struct1 = self.parser.get_structure("1lcd", "PDB/1LCD.pdb") 290 291 confirm_numbering(struct1) 292 293 # Round trip: serialize and parse again 294 self.io.set_structure(struct1) 295 filenumber, filename = tempfile.mkstemp() 296 os.close(filenumber) 297 try: 298 self.io.save(filename) 299 struct2 = self.parser.get_structure("1lcd", filename) 300 confirm_numbering(struct2) 301 confirm_single_end(filename) 302 finally: 303 os.remove(filename) 304 305 def test_pdbio_write_x_element(self): 306 """Write a structure with atomic element X with PDBIO.""" 307 struct1 = self.structure 308 309 # Change element of one atom 310 atom = next(struct1.get_atoms()) 311 atom.element = "X" # X is assigned in Atom.py as last resort 312 313 self.io.set_structure(struct1) 314 315 filenumber, filename = tempfile.mkstemp() 316 os.close(filenumber) 317 318 try: 319 self.io.save(filename) 320 finally: 321 os.remove(filename) 322 323 def test_pdbio_write_unk_element(self): 324 """PDBIO raises ValueError when writing unrecognised atomic elements.""" 325 struct1 = self.structure 326 327 atom = next(struct1.get_atoms()) 328 atom.element = "1" 329 330 self.io.set_structure(struct1) 331 332 filenumber, filename = tempfile.mkstemp() 333 os.close(filenumber) 334 335 with self.assertRaises(ValueError): 336 self.io.save(filename) 337 os.remove(filename) 338 339 340if __name__ == "__main__": 341 runner = unittest.TextTestRunner(verbosity=2) 342 unittest.main(testRunner=runner) 343