Make vcf classes immutable so they can safely be used as keys.

Author: donkirkby

pyvdrm/tests/test_vcf.py

@@ -77,6 +77,12 @@ def test_equal_with_wildtype_mismatch(self):
             if set1 == set2:
                 pass
 
+    def test_immutable(self):
+        m = Mutation('Q1A')
+
+        with self.assertRaises(AttributeError):
+            m.pos = 2
+
 
 class TestMutationSet(unittest.TestCase):
     def test_init_text(self):
@@ -251,6 +257,12 @@ def test_no_wildtype(self):
             mutations = MutationSet(text)
             self.assertIsNone(mutations.wildtype)
 
+    def test_immutable(self):
+        ms = MutationSet('Q80KR')
+
+        with self.assertRaises(AttributeError):
+            ms.wildtype = 'D'
+
 
 class TestVariantCalls(unittest.TestCase):
     def test_init_text(self):
@@ -264,22 +276,21 @@ def test_init_text(self):
     def test_init_single_sequence(self):
         reference = 'ACHE'
         sample = 'ICRE'
-        expected_mutation_sets = {MutationSet('A1I'), MutationSet('H3R')}
+        expected_calls = VariantCalls('A1I C2C H3R E4E')
 
         calls = VariantCalls(reference=reference, sample=sample)
 
         self.assertEqual(reference, calls.reference)
-        self.assertEqual(expected_mutation_sets, calls.mutation_sets)
+        self.assertEqual(expected_calls, calls)
 
     def test_init_multiple_sequences(self):
         reference = 'ACHE'
         sample = ['IN', 'C', 'HR', 'E']
-        expected_mutation_sets = {MutationSet('A1IN'), MutationSet('H3HR')}
+        expected_calls = VariantCalls('A1IN C2C H3HR E4E')
 
         calls = VariantCalls(reference=reference, sample=sample)
 
-        self.assertEqual(reference, calls.reference)
-        self.assertEqual(expected_mutation_sets, calls.mutation_sets)
+        self.assertEqual(expected_calls, calls)
 
     def test_init_bad_length(self):
         reference = 'ACHE'
@@ -289,6 +300,16 @@ def test_init_bad_length(self):
                 ValueError, r'Reference length was 4 and sample length was 5\.'):
             VariantCalls(reference=reference, sample=sample)
 
+    def test_init_with_reference(self):
+        expected_reference = 'ASH'
+        expected_repr = "VariantCalls('A1IL H3R')"
+
+        calls = VariantCalls('1IL 3R', reference=expected_reference)
+        r = repr(calls)
+
+        self.assertEqual(expected_reference, calls.reference)
+        self.assertEqual(expected_repr, r)
+
     def test_repr(self):
         expected_repr = "VariantCalls('A1IL H3R')"
         calls = VariantCalls('A1IL H3R')
@@ -325,6 +346,28 @@ def test_hash(self):
         self.assertNotEqual(hash1, hash3)
         self.assertNotEqual(hash1, hash4)
 
+    def test_iter(self):
+        calls = VariantCalls('A1IL H3R')
+        expected_mutation_sets = {MutationSet('A1IL'), MutationSet('H3R')}
+
+        mutation_sets = set(calls)
+
+        self.assertEqual(expected_mutation_sets, mutation_sets)
+
+    def test_in(self):
+        calls = VariantCalls('A1IL H3R')
+        mutation_set1 = MutationSet('H3R')
+        mutation_set2 = MutationSet('H4R')
+
+        self.assertIn(mutation_set1, calls)
+        self.assertNotIn(mutation_set2, calls)
+
+    def test_immutable(self):
+        calls = VariantCalls('A1IL H3R')
+
+        with self.assertRaises(AttributeError):
+            calls.reference = 'ASH'
+
 
 if __name__ == '__main__':
     unittest.main()

pyvdrm/vcf.py

@@ -2,32 +2,45 @@
 Classes for dealing with amino acid mutation sets
 """
 import re
+from collections import namedtuple
 from operator import attrgetter
 
 
-class VariantCalls:
+class VariantCalls(namedtuple('VariantCalls', 'mutation_sets reference')):
+    # TODO: remove all these __init__ methods once PyCharm bug is fixed.
+    # https://youtrack.jetbrains.com/issue/PY-26834
+    # noinspection PyUnusedLocal
     def __init__(self, text=None, reference=None, sample=None):
         """ Construct a set of Mutations given two aligned amino acid sequences
 
         :param str reference: the wild-type reference
         :param sample: amino acids present at each position, either a string or
         a list of strings
         """
-        self.reference = reference
+        # noinspection PyArgumentList
+        super().__init__()
+
+    def __new__(cls, text=None, reference=None, sample=None):
         if text is not None:
             terms = text.split()
-            self.mutation_sets = frozenset(MutationSet(term) for term in terms)
+            mutation_sets = frozenset(
+                MutationSet(term, reference=reference)
+                for term in terms)
         else:
             if len(reference) != len(sample):
                 raise ValueError(
                     'Reference length was {} and sample length was {}.'.format(
                         len(reference),
                         len(sample)))
 
-            self.mutation_sets = {MutationSet(pos=i, variants=alt, wildtype=ref)
-                                  for i, (alt, ref) in enumerate(zip(sample, reference),
-                                                                 1)
-                                  if ref != alt}
+            mutation_sets = {MutationSet(pos=i, variants=alt, wildtype=ref)
+                             for i, (alt, ref) in enumerate(zip(sample,
+                                                                reference),
+                                                            1)}
+        # noinspection PyArgumentList
+        return super().__new__(cls,
+                               mutation_sets=mutation_sets,
+                               reference=reference)
 
     def __str__(self):
         return ' '.join(map(str, sorted(self.mutation_sets,
@@ -40,23 +53,23 @@ def __repr__(self):
     def __eq__(self, other):
         return self.mutation_sets == other.mutation_sets
 
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def __hash__(self):
         return hash(self.mutation_sets)
 
+    def __iter__(self):
+        return iter(self.mutation_sets)
 
-class Mutation(object):
-    """Mutation has optional wildtype, position, and call"""
+    def __contains__(self, item):
+        return item in self.mutation_sets
 
-    def __init__(self, text=None, wildtype=None, pos=None, variant=None):
-        """ Initialize.
 
-        :param str text: will be parsed for wildtype (optional), position,
-            and variant
-        :param str wildtype: amino acid abbreviation for wild type
-        :param str|int pos: position
-        :param str variant: single amino acid abbreviation, or 'i' for
-            insertion, or 'd' for deletion
-        """
+class Mutation(namedtuple('Mutation', 'pos variant wildtype')):
+    """Mutation has optional wildtype, position, and call"""
+
+    def __new__(cls, text=None, wildtype=None, pos=None, variant=None):
         if text is not None:
             match = re.match(r"([A-Z]?)(\d+)([idA-Z])", text)
             if match is None:
@@ -68,13 +81,25 @@ def __init__(self, text=None, wildtype=None, pos=None, variant=None):
                 raise ValueError('Mutation text only allows one variant.')
 
             wildtype, pos, variant = match.groups()
-        self.wildtype = wildtype or None
-        self.pos = int(pos)
-        self.variant = variant
+        # noinspection PyArgumentList
+        return super().__new__(cls,
+                               pos=int(pos),
+                               variant=variant,
+                               wildtype=wildtype or None)
 
-    def extract_wildtype(self, seq):
-        """I really don't like this; please don't actually use this"""
-        self.wildtype = seq[self.pos - 1]
+    # noinspection PyUnusedLocal
+    def __init__(self, text=None, wildtype=None, pos=None, variant=None):
+        """ Initialize.
+
+        :param str text: will be parsed for wildtype (optional), position,
+            and variant
+        :param str wildtype: amino acid abbreviation for wild type
+        :param str|int pos: position
+        :param str variant: single amino acid abbreviation, or 'i' for
+            insertion, or 'd' for deletion
+        """
+        # noinspection PyArgumentList
+        super().__init__()
 
     def __repr__(self):
         text = str(self)
@@ -100,36 +125,31 @@ def __eq__(self, other):
         # now that we agree on the wt and position
         return (self.pos, self.variant) == (other.pos, other.variant)
 
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def __hash__(self):
         return hash((self.pos, self.variant))
 
 
-class MutationSet(object):
+class MutationSet(namedtuple('MutationSet', 'pos mutations wildtype')):
     """Handle sets of mutations at a position"""
 
-    def __init__(self,
-                 text=None,
-                 wildtype=None,
-                 pos=None,
-                 variants=None,
-                 mutations=None):
-        """ Initialize
-
-        :param str text: will be parsed for wildtype (optional), position,
-            and variants
-        :param str wildtype: amino acid abbreviation for wild type
-        :param int|str pos: position
-        :param str variants: zero or more amino acid abbreviations, or 'i' for
-            insertion, or 'd' for deletion
-        :param mutations: a sequence of Mutation objects, with matching
-            positions and wild types
-        """
+    def __new__(cls,
+                text=None,
+                wildtype=None,
+                pos=None,
+                variants=None,
+                mutations=None,
+                reference=None):
         if text:
             match = re.match(r"([A-Z]?)(\d+)([idA-Z]*)", text)
             if match is None:
                 raise ValueError
 
             wildtype, pos, variants = match.groups()
+            if reference:
+                wildtype = reference[int(pos)-1]
 
         if variants:
             mutations = frozenset(Mutation(wildtype=wildtype,
@@ -159,18 +179,41 @@ def __init__(self,
             pos = positions.pop()
             if wildtypes:
                 wildtype = wildtypes.pop()
-        self.wildtype = wildtype or None
-        self.pos = int(pos)
-        self.mutations = mutations
+        # noinspection PyArgumentList
+        return super().__new__(cls,
+                               wildtype=wildtype or None,
+                               pos=int(pos),
+                               mutations=mutations)
+
+    # noinspection PyUnusedLocal
+    def __init__(self,
+                 text=None,
+                 wildtype=None,
+                 pos=None,
+                 variants=None,
+                 mutations=None,
+                 reference=None):
+        """ Initialize
+
+        :param str text: will be parsed for wildtype (optional), position,
+            and variants
+        :param str wildtype: amino acid abbreviation for wild type
+        :param int|str pos: position
+        :param str variants: zero or more amino acid abbreviations, or 'i' for
+            insertion, or 'd' for deletion
+        :param mutations: a sequence of Mutation objects, with matching
+            positions and wild types
+        :param str reference: alternative source for wildtype, based on
+            pos - 1
+        """
+        # noinspection PyArgumentList
+        super().__init__()
 
     def __len__(self):
         return len(self.mutations)
 
     def __contains__(self, call):
-        for mutation in self.mutations:
-            if call == mutation:
-                return True
-        return False
+        return call in self.mutations
 
     def __eq__(self, other):
         if self.pos != other.pos:
@@ -184,18 +227,14 @@ def __eq__(self, other):
                 raise ValueError(message)
         return self.mutations == other.mutations
 
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def __hash__(self):
         return hash((self.pos, self.mutations))
 
     def __iter__(self):
-        self._mu_iter = list(self.mutations).__iter__()
-        return self._mu_iter
-
-    def __next__(self):
-        return self._mu_iter.__next__()
-
-    def __reversed__(self):
-        return list(self.mutations).__reversed__()
+        return iter(self.mutations)
 
     def __str__(self):
         text = self.wildtype or ''