From 67b1ce11f33565f6eedfdd1bfd334329bc4f19e1 Mon Sep 17 00:00:00 2001 From: joaomcteixeira Date: Tue, 12 Dec 2023 20:38:02 +0100 Subject: [PATCH] solves #160 --- pdbtools/pdb_tidy.py | 12 ++++----- tests/data/hetatm_ensemble.pdb | 45 ++++++++++++++++++++++++++++++++++ tests/test_pdb_tidy.py | 27 ++++++++++++++++++++ 3 files changed, 78 insertions(+), 6 deletions(-) create mode 100644 tests/data/hetatm_ensemble.pdb diff --git a/pdbtools/pdb_tidy.py b/pdbtools/pdb_tidy.py index fc39df55..39bde54b 100644 --- a/pdbtools/pdb_tidy.py +++ b/pdbtools/pdb_tidy.py @@ -214,9 +214,9 @@ def make_TER(prev_line): if atom_section: atom_section = False yield make_TER(prev_line) - if in_model: - yield "{:<80}\n".format("ENDMDL") - in_model = False + if in_model: + yield "{:<80}\n".format("ENDMDL") + in_model = False if line.startswith('MODEL'): line = "MODEL " + " " + str(num_models).rjust(4) @@ -240,9 +240,9 @@ def make_TER(prev_line): # Add last TER statement atom_section = False yield make_TER(prev_line) - if in_model: - yield "{:<80}\n".format("ENDMDL") - in_model = False + if in_model: + yield "{:<80}\n".format("ENDMDL") + in_model = False # Add END statement yield "{:<80}\n".format("END") diff --git a/tests/data/hetatm_ensemble.pdb b/tests/data/hetatm_ensemble.pdb new file mode 100644 index 00000000..ca4707df --- /dev/null +++ b/tests/data/hetatm_ensemble.pdb @@ -0,0 +1,45 @@ +MODEL 1 +HETATM 1 C1 G39 B 500 -30.374 -53.143 6.122 1.00 31.18 C +HETATM 2 O1A G39 B 500 -31.439 -52.533 5.850 1.00 30.30 O +HETATM 3 O1B G39 B 500 -29.257 -52.625 5.879 1.00 31.23 O +HETATM 4 C2 G39 B 500 -30.437 -54.486 6.725 1.00 30.48 C +HETATM 5 C3 G39 B 500 -29.159 -55.188 7.132 1.00 30.04 C +HETATM 6 C4 G39 B 500 -29.380 -56.162 8.294 1.00 30.37 C +HETATM 7 N4 G39 B 500 -28.195 -56.994 8.430 1.00 30.09 N1+ +HETATM 8 C5 G39 B 500 -30.627 -57.033 8.097 1.00 29.53 C +HETATM 9 N5 G39 B 500 -30.803 -57.955 9.206 1.00 28.16 N +HETATM 10 C6 G39 B 500 -31.883 -56.178 7.937 1.00 30.53 C +HETATM 11 C7 G39 B 500 -31.632 -55.093 6.901 1.00 30.75 C +HETATM 12 O7 G39 B 500 -33.007 -56.989 7.531 1.00 31.80 O +HETATM 13 C8 G39 B 500 -34.158 -56.905 8.398 1.00 30.52 C +HETATM 14 C9 G39 B 500 -34.879 -55.553 8.283 1.00 31.36 C +HETATM 15 C10 G39 B 500 -30.683 -59.283 9.087 1.00 27.13 C +HETATM 16 O10 G39 B 500 -30.629 -59.863 8.010 1.00 26.14 O +HETATM 17 C11 G39 B 500 -30.611 -60.040 10.380 1.00 26.05 C +HETATM 18 C81 G39 B 500 -35.153 -58.020 8.079 1.00 29.88 C +HETATM 19 C82 G39 B 500 -34.654 -59.395 8.467 1.00 28.23 C +HETATM 20 C91 G39 B 500 -35.425 -55.279 6.891 1.00 32.33 C +ENDMDL +MODEL 2 +HETATM 1 C1 G39 B 500 -30.374 -53.143 6.122 1.00 31.18 C +HETATM 2 O1A G39 B 500 -31.439 -52.533 5.850 1.00 30.30 O +HETATM 3 O1B G39 B 500 -29.257 -52.625 5.879 1.00 31.23 O +HETATM 4 C2 G39 B 500 -30.437 -54.486 6.725 1.00 30.48 C +HETATM 5 C3 G39 B 500 -29.159 -55.188 7.132 1.00 30.04 C +HETATM 6 C4 G39 B 500 -29.380 -56.162 8.294 1.00 30.37 C +HETATM 7 N4 G39 B 500 -28.195 -56.994 8.430 1.00 30.09 N1+ +HETATM 8 C5 G39 B 500 -30.627 -57.033 8.097 1.00 29.53 C +HETATM 9 N5 G39 B 500 -30.803 -57.955 9.206 1.00 28.16 N +HETATM 10 C6 G39 B 500 -31.883 -56.178 7.937 1.00 30.53 C +HETATM 11 C7 G39 B 500 -31.632 -55.093 6.901 1.00 30.75 C +HETATM 12 O7 G39 B 500 -33.007 -56.989 7.531 1.00 31.80 O +HETATM 13 C8 G39 B 500 -34.158 -56.905 8.398 1.00 30.52 C +HETATM 14 C9 G39 B 500 -34.879 -55.553 8.283 1.00 31.36 C +HETATM 15 C10 G39 B 500 -30.683 -59.283 9.087 1.00 27.13 C +HETATM 16 O10 G39 B 500 -30.629 -59.863 8.010 1.00 26.14 O +HETATM 17 C11 G39 B 500 -30.611 -60.040 10.380 1.00 26.05 C +HETATM 18 C81 G39 B 500 -35.153 -58.020 8.079 1.00 29.88 C +HETATM 19 C82 G39 B 500 -34.654 -59.395 8.467 1.00 28.23 C +HETATM 20 C91 G39 B 500 -35.425 -55.279 6.891 1.00 32.33 C +ENDMDL +END \ No newline at end of file diff --git a/tests/test_pdb_tidy.py b/tests/test_pdb_tidy.py index e3cbec2f..42d96ec9 100644 --- a/tests/test_pdb_tidy.py +++ b/tests/test_pdb_tidy.py @@ -260,6 +260,33 @@ def test_corrects_model_ENDMDL(self): self.assertEqual(self.stdout[-1].strip(), "END") self.assertEqual(len(self.stdout[-1]), 80) + def test_corrects_model_ENDMDL_with_HETATM(self): + """Correct MODEL lines.""" + fpath = os.path.join(data_dir, 'hetatm_ensemble.pdb') + sys.argv = ['', fpath] + self.exec_module() + self.assertEqual(self.retcode, 0) + self.assertEqual(len(self.stdout), 45) + self.assertEqual(len(self.stderr), 0) + + # MODEL lines + self.assertEqual(len(self.stdout[0]), 80) + self.assertEqual(len(self.stdout[0].strip()), 14) + self.assertEqual(len(self.stdout[22]), 80) + self.assertEqual(len(self.stdout[22].strip()), 14) + self.assertEqual(self.stdout[0].strip(), "MODEL 1") + self.assertEqual(self.stdout[22].strip(), "MODEL 2") + + # ENDMDL LINES + self.assertEqual(len(self.stdout[21]), 80) + self.assertEqual(len(self.stdout[43]), 80) + self.assertEqual(self.stdout[21].strip(), "ENDMDL") + self.assertEqual(self.stdout[43].strip(), "ENDMDL") + + # END LINES + self.assertEqual(self.stdout[-1].strip(), "END") + self.assertEqual(len(self.stdout[-1]), 80) + def test_file_not_found(self): """$ pdb_tidy not_existing.pdb"""