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