Add residue codes

samirelanduk · May 17, 2018 · e676eb2 · e676eb2
1 parent 1c75fb1
commit e676eb2
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 0 deletions.
diff --git a/atomium/structures/molecules.py b/atomium/structures/molecules.py
@@ -745,6 +745,17 @@ def full_name(self):
             return RESIDUES.get(self._name.upper(), self._name)
 
 
+    @property
+    def code(self):
+        """Returns the amino acid code of the reside if it is one of the 20
+        canonical amino acids - otherwise it just returns 'X'.
+
+        :rtype: ``str``"""
+
+        if self._name:
+            return CODES.get(self._name.upper(), "X")
+
+
     @property
     def next(self):
         """Residues can be linked to each other in a linear chain. This property
@@ -832,3 +843,9 @@ def chain(self):
  "ASP": "aspartic acid", "GLU": "glutamic acid", "LYS": "lysine",
  "ARG": "arginine", "HIS": "histidine", "HOH": "water"
 }
+
+CODES = {
+ "VAL":"V", "ILE":"I", "LEU":"L", "GLU":"E", "GLN":"Q", "ASP":"D", "ASN":"N",
+ "HIS":"H", "TRP":"W", "PHE":"F", "TYR":"Y", "ARG":"R", "LYS":"K", "SER":"S",
+ "THR":"T", "MET":"M", "ALA":"A", "GLY":"G", "PRO":"P", "CYS":"C"
+}
diff --git a/tests/integration/test_structures.py b/tests/integration/test_structures.py
@@ -264,8 +264,12 @@ def test_atoms_in_molecules(self):
         self.assertEqual(res1.full_name, "valine")
         self.assertEqual(res2.full_name, "glutamic acid")
         self.assertEqual(res3.full_name, "tryptophan")
+        self.assertEqual(res1.code, "V")
+        self.assertEqual(res2.code, "E")
+        self.assertEqual(res3.code, "W")
         res3.name = "XYZ"
         self.assertEqual(res3.full_name, "XYZ")
+        self.assertEqual(res3.code, "X")
 
         # Copies
         copy = molecule.copy()

diff --git a/tests/unit/structure_tests/test_residues.py b/tests/unit/structure_tests/test_residues.py
@@ -79,6 +79,22 @@ def test_can_lookup_full_name(self):
 
 
 
+class ResidueCodeTests(ResidueTest):
+
+    def test_can_get_generic_code(self):
+        res = Residue(self.atom1, self.atom2, self.atom3)
+        self.assertIsNone(res.code)
+        res._name = "XMP"
+        self.assertEqual(res.code, "X")
+
+
+    def test_can_lookup_code(self):
+        res = Residue(self.atom1, self.atom2, self.atom3)
+        res._name = "met"
+        self.assertEqual(res.code, "M")
+
+
+
 class ResidueNextTests(ResidueTest):
 
     def test_next_gets_next(self):