Repository: acorg/seqgen

Issue:
Add a 'name count' directive so a certain sequence in a spec can be named
Currently we only assign the name of a named spec to the first sequence in that series. Make it so any one can be named.

Discussion:

Patch:
diff --git a/README.md b/README.md
index af241af..40f7376 100644
--- a/README.md
+++ b/README.md
@@ -146,6 +146,10 @@ specification keys you can put in a object in the JSON is as follows:
 * `name`: Give a name to the sequence so you can refer to it from elsewhere
   in the JSON. Note that this has nothing to do with FASTA ids, it is just
   an internal name.
+* `name count`: An optional integer specifying which of a set of sequences
+  generated from the same specification should be named (assumes `name`
+  is also given). The default is to name the first sequence. Note that this
+  number is 1-based. Special values `first` and `last` are also allowed.
 * `random aa`: The sequence should be made of random amino acids.
 * `random nt`: The sequence should be made of random nucleotides.
 * `ratchet`: If `true` and a count and mutation rate are given, sequences
diff --git a/seqgen/__init__.py b/seqgen/__init__.py
index 285814d..51d6e92 100644
--- a/seqgen/__init__.py
+++ b/seqgen/__init__.py
@@ -2,6 +2,6 @@ from seqgen.sequences import Sequences
 
 # Note that the version string must have the following format, otherwise it
 # will not be found by the version() function in ../setup.py
-__version__ = '1.0.11'
+__version__ = '1.0.12'
 
 __all__ = ['Sequences']
diff --git a/seqgen/sequences.py b/seqgen/sequences.py
index 3af0d90..b57ba6d 100644
--- a/seqgen/sequences.py
+++ b/seqgen/sequences.py
@@ -32,6 +32,7 @@ class Sequences(object):
         'length',
         'mutation rate',
         'name',
+        'name count',
         'random aa',
         'random nt',
         'ratchet',
@@ -88,6 +89,27 @@ class Sequences(object):
         self._checkValid()
         self._names = {}
 
+    def _nameCount(self, spec):
+        """
+        Get the name count for a spec
+
+        @param spec: A C{dict} with keys/values specifying a sequence.
+        @return: A 0-based integer count of the sequence index that should be
+            named or C{0} if there is no 'name count' directive in C{spec}.
+        """
+        nSequences = spec.get('count', 1)
+        try:
+            nameCount = spec['name count']
+        except KeyError:
+            return 0
+        else:
+            if nameCount == 'first':
+                return 0
+            elif nameCount == 'last':
+                return nSequences - 1
+            else:
+                return nameCount - 1
+
     def _checkValid(self):
         """
         Check that all specification dicts contain sensible values.
@@ -96,6 +118,7 @@ class Sequences(object):
             to be produced.
         @raise ValueError: If any problem is found.
         """
+        names = set()
         for specCount, spec in enumerate(self._sequenceSpecs, start=1):
             if spec.get('ratchet'):
                 nSequences = spec.get('count', 1)
@@ -109,6 +132,27 @@ class Sequences(object):
                         'Sequence specification %d is specified as ratchet '
                         'but does not give a mutation rate.' % specCount)
 
+            nSequences = spec.get('count', 1)
+
+            nameCount = self._nameCount(spec)
+            if nameCount is not None and nameCount >= nSequences:
+                raise ValueError(
+                    'Sequence specification %d will only produce %d '
+                    'sequence%s but has a (too high) name count value of %d.' %
+                    (specCount, nSequences, '' if nSequences == 1 else 's',
+                     nameCount + 1))
+
+            try:
+                name = spec['name']
+            except KeyError:
+                pass
+            else:
+                if name in names:
+                    raise ValueError("Name '%s' is duplicated in the JSON "
+                                     "specification." % name)
+                else:
+                    names.add(name)
+
     def _checkKeys(self):
         """
         Check that all specification dicts only contain legal keys.
@@ -286,6 +330,8 @@ class Sequences(object):
         alphabet = None
         previousRead = None
         nSequences = spec.get('count', 1)
+        nameCount = self._nameCount(spec)
+
         for count in range(nSequences):
             id_ = None
             if 'sections' in spec:
@@ -329,14 +375,10 @@ class Sequences(object):
             read = Read(id_, sequence)
             read.alphabet = alphabet
 
-            # Keep a reference to this result if it is named.
-            if count == 0 and 'name' in spec:
-                name = spec['name']
-                if name in self._names:
-                    raise ValueError("Name '%s' is duplicated in the JSON "
-                                     "specification." % name)
-                else:
-                    self._names[name] = read
+            # Keep a reference to this result if it is named and this is the
+            # sequence (by count) we want to name.
+            if 'name' in spec and nameCount == count:
+                self._names[spec['name']] = read
 
             if not spec.get('skip'):
                 yield read


Test patch:
diff --git a/test/testSequences.py b/test/testSequences.py
index b4b839d..a57de94 100644
--- a/test/testSequences.py
+++ b/test/testSequences.py
@@ -56,9 +56,8 @@ class TestSequences(TestCase):
         If a duplicate sequence name is present in the JSON, a ValueError
         must be raised.
         """
-        s = Sequences(spec='filename')
         error = "^Name 'a' is duplicated in the JSON specification\\.$"
-        assertRaisesRegex(self, ValueError, error, list, s)
+        assertRaisesRegex(self, ValueError, error, Sequences, spec='filename')
 
     def testNoSequences(self):
         """
@@ -738,3 +737,138 @@ class TestSequences(TestCase):
         # The sequences of the original and the second mutant must be
         # identical.
         self.assertEqual(orig.sequence, mutant2.sequence)
+
+    def testNameCountTooHigh(self):
+        """
+        A name count value that is higher than the number of sequences in
+        a specification must result in a ValueError.
+        """
+        spec = StringIO('''{
+            "sequences": [
+                {
+                    "name": "orig",
+                    "name count": 10,
+                    "count": 2
+                }
+            ]
+        }''')
+        error = ('^Sequence specification 1 will only produce 2 sequences but '
+                 'has a \\(too high\\) name count value of 10\\.$')
+        assertRaisesRegex(self, ValueError, error, Sequences, spec)
+
+    def testNameCountFirst(self):
+        """
+        A name count value of 'first' must result in the first sequence of a
+        series being named.
+        """
+        s = Sequences(StringIO('''{
+            "sequences": [
+                {
+                    "name": "orig",
+                    "name count": "first",
+                    "count": 20,
+                    "mutation rate": 0.1,
+                    "length": 500,
+                    "ratchet": true
+                },
+                {
+                    "id": "copy",
+                    "from name": "orig"
+                }
+            ]
+        }'''))
+        reads = list(s)
+        copySequence = reads[-1].sequence
+        # The final sequence (the copy) must be the same as the first sequence
+        # (and no others).
+        self.assertEqual(reads[0].sequence, copySequence)
+        for i in range(1, 20):
+            self.assertNotEqual(reads[i].sequence, copySequence)
+
+    def testNameCountLast(self):
+        """
+        A name count value of 'last' must result in the last sequence of a
+        series being named.
+        """
+        s = Sequences(StringIO('''{
+            "sequences": [
+                {
+                    "name": "orig",
+                    "name count": "last",
+                    "count": 20,
+                    "mutation rate": 0.1,
+                    "length": 500,
+                    "ratchet": true
+                },
+                {
+                    "id": "copy",
+                    "from name": "orig"
+                }
+            ]
+        }'''))
+        reads = list(s)
+        copySequence = reads[-1].sequence
+        # The final sequence (the copy) must be the same as the second last
+        # (i.e., the final mutant) sequence (and no others).
+        self.assertEqual(reads[-2].sequence, copySequence)
+        for i in range(19):
+            self.assertNotEqual(reads[i].sequence, copySequence)
+
+    def testNameCountOne(self):
+        """
+        A name count value of 1 must result in the first sequence of a
+        series being named.
+        """
+        s = Sequences(StringIO('''{
+            "sequences": [
+                {
+                    "name": "orig",
+                    "name count": 1,
+                    "count": 20,
+                    "mutation rate": 0.1,
+                    "length": 500,
+                    "ratchet": true
+                },
+                {
+                    "id": "copy",
+                    "from name": "orig"
+                }
+            ]
+        }'''))
+        reads = list(s)
+        copySequence = reads[-1].sequence
+        # The final sequence (the copy) must be the same as the first sequence
+        # (and no others).
+        self.assertEqual(reads[0].sequence, copySequence)
+        for i in range(1, 20):
+            self.assertNotEqual(reads[i].sequence, copySequence)
+
+    def testNameCountTwo(self):
+        """
+        A name count value of 2 must result in the second sequence of a
+        series being named.
+        """
+        s = Sequences(StringIO('''{
+            "sequences": [
+                {
+                    "name": "orig",
+                    "name count": 2,
+                    "count": 20,
+                    "mutation rate": 0.1,
+                    "length": 500,
+                    "ratchet": true
+                },
+                {
+                    "id": "copy",
+                    "from name": "orig"
+                }
+            ]
+        }'''))
+        reads = list(s)
+        copySequence = reads[-1].sequence
+        # The final sequence (the copy) must be the same as the second sequence
+        # (and no others).
+        self.assertEqual(reads[1].sequence, copySequence)
+        self.assertNotEqual(reads[0].sequence, copySequence)
+        for i in range(2, 20):
+            self.assertNotEqual(reads[i].sequence, copySequence)