Repository: acorg/dark-matter

Issue:
Add an option to the `PaddedSAM.queries` method to have it also return the pysam.AlignedSegment

Comments:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 67af944..2065375 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,11 @@
+## 3.0.28 July 3, 2018
+
+Added `alsoYieldAlignments` option to `PaddedSAM.queries` method to have
+the returned generator also yield the `pysam.AlignedSegment` instance with
+the gap-padded query sequence. This makes it possible to retrieve padded
+queries from SAM/BAM and generate SAM/BAM (or FASTQ) of some subset of the
+queries.
+
 ## 3.0.27 June 30, 2018
 
 Added `bin/filter-sam.py` script.
diff --git a/bin/filter-sam.py b/bin/filter-sam.py
index 18b5b47..363fed1 100755
--- a/bin/filter-sam.py
+++ b/bin/filter-sam.py
@@ -38,26 +38,26 @@ if __name__ == '__main__':
               'error message is printed unless --quiet is used).'))
 
     parser.add_argument(
-         '--dropUnmapped', default=False, action='store_true',
-         help='If given, unmapped matches will not be output.')
+        '--dropUnmapped', default=False, action='store_true',
+        help='If given, unmapped matches will not be output.')
 
     parser.add_argument(
-         '--dropSecondary', default=False, action='store_true',
-         help='If given, secondary matches will not be output.')
+        '--dropSecondary', default=False, action='store_true',
+        help='If given, secondary matches will not be output.')
 
     parser.add_argument(
-         '--dropSupplementary', default=False, action='store_true',
-         help='If given, supplementary matches will not be output.')
+        '--dropSupplementary', default=False, action='store_true',
+        help='If given, supplementary matches will not be output.')
 
     parser.add_argument(
-         '--dropDuplicates', default=False, action='store_true',
-         help=('If given, matches flagged as optical or PCR duplicates will '
-               'not be output.'))
+        '--dropDuplicates', default=False, action='store_true',
+        help=('If given, matches flagged as optical or PCR duplicates will '
+              'not be output.'))
 
     parser.add_argument(
-         '--keepQCFailures', default=False, action='store_true',
-         help=('If given, reads that are considered quality control failures '
-               'will be included in the output.'))
+        '--keepQCFailures', default=False, action='store_true',
+        help=('If given, reads that are considered quality control failures '
+              'will be included in the output.'))
 
     parser.add_argument(
         '--referenceWhitelist', metavar='NAME', action='append',
@@ -106,17 +106,17 @@ if __name__ == '__main__':
 
             if (filterRead(Read(alignment.query_name,
                                 alignment.query_sequence,
-                                alignment.qual))
-                    and not (
-                        (alignment.is_unmapped and dropUnmapped) or
-                        (alignment.is_secondary and dropSecondary) or
-                        (alignment.is_supplementary and dropSupplementary) or
-                        (alignment.is_duplicate and dropDuplicates) or
-                        (alignment.is_qcfail and not keepQCFailures) or
-                        (referenceWhitelist is not None and
-                         alignment.reference_name not in referenceWhitelist) or
-                        (referenceBlacklist is not None and
-                         alignment.reference_name in referenceBlacklist))):
+                                alignment.qual)) and
+                not (
+                    (alignment.is_unmapped and dropUnmapped) or
+                    (alignment.is_secondary and dropSecondary) or
+                    (alignment.is_supplementary and dropSupplementary) or
+                    (alignment.is_duplicate and dropDuplicates) or
+                    (alignment.is_qcfail and not keepQCFailures) or
+                    (referenceWhitelist is not None and
+                     alignment.reference_name not in referenceWhitelist) or
+                    (referenceBlacklist is not None and
+                     alignment.reference_name in referenceBlacklist))):
                 kept += 1
                 save(alignment)
         out.close()
diff --git a/dark/__init__.py b/dark/__init__.py
index dabfdd7..2f8e41d 100644
--- a/dark/__init__.py
+++ b/dark/__init__.py
@@ -7,4 +7,4 @@ if sys.version_info < (2, 7):
 # will not be found by the version() function in ../setup.py
 #
 # Remember to update ../CHANGELOG.md describing what's new in each version.
-__version__ = '3.0.27'
+__version__ = '3.0.28'
diff --git a/dark/sam.py b/dark/sam.py
index f2bb28b..a5d9363 100644
--- a/dark/sam.py
+++ b/dark/sam.py
@@ -61,7 +61,7 @@ class PaddedSAM(object):
                 dropSecondary=False, dropSupplementary=False,
                 dropDuplicates=False, allowDuplicateIds=False,
                 keepQCFailures=False, rcNeeded=False, padChar='-',
-                queryInsertionChar='N'):
+                queryInsertionChar='N', alsoYieldAlignments=False):
         """
         Produce padded (with gaps) queries according to the CIGAR string and
         reference sequence length for each matching query sequence.
@@ -107,13 +107,18 @@ class PaddedSAM(object):
             is inserted as a 'missing' query character (i.e., a base that can
             be assumed to have been lost due to an error) whose existence is
             necessary for the match to continue.
+        @param alsoYieldAlignments: If C{True} the returned generator will
+            yield 2-tuples containing a padded query and the
+            C{pysam.AlignedSegment} for each query.
         @raises UnequalReferenceLengthError: If C{referenceName} is C{None}
             and the reference sequence lengths in the SAM/BAM file are not all
             identical.
         @raises UnknownReference: If C{referenceName} does not exist.
         @return: A generator that yields C{Read} instances that are padded
             with gap characters to align them to the length of the reference
-            sequence.
+            sequence. See C{alsoYieldAlignments}, above, to have the generator
+            yield tuples also containing the corresponding
+            C{pysam.AlignedSegment}.
         """
         samfile = self.samfile
 
@@ -142,24 +147,24 @@ class PaddedSAM(object):
         MATCH_OPERATIONS = {CMATCH, CEQUAL, CDIFF}
         lastQuery = None
 
-        for lineNumber, read in enumerate(samfile.fetch(), start=1):
-            if (read.is_unmapped or
-                    (read.is_secondary and dropSecondary) or
-                    (read.is_supplementary and dropSupplementary) or
-                    (read.is_duplicate and dropDuplicates) or
-                    (read.is_qcfail and not keepQCFailures) or
+        for lineNumber, alignment in enumerate(samfile.fetch(), start=1):
+            if (alignment.is_unmapped or
+                    (alignment.is_secondary and dropSecondary) or
+                    (alignment.is_supplementary and dropSupplementary) or
+                    (alignment.is_duplicate and dropDuplicates) or
+                    (alignment.is_qcfail and not keepQCFailures) or
                     (referenceId is not None and
-                     read.reference_id != referenceId)):
+                     alignment.reference_id != referenceId)):
                 continue
 
-            query = read.query_sequence
+            query = alignment.query_sequence
 
             # Secondary (and presumably supplementary) alignments may have
             # a '*' (None in pysam) SEQ field, indicating that the previous
             # sequence should be used. This is best practice according to
             # section 2.5.2 of https://samtools.github.io/hts-specs/SAMv1.pdf
             if query is None:
-                if read.is_secondary or read.is_supplementary:
+                if alignment.is_secondary or alignment.is_supplementary:
                     if lastQuery is None:
                         raise ValueError(
                             'Query line %d has an empty SEQ field, but no '
@@ -176,29 +181,29 @@ class PaddedSAM(object):
             # due to it being reverse complimented for the alignment).
             lastQuery = query
 
-            if read.is_reverse:
+            if alignment.is_reverse:
                 if rcNeeded:
                     query = DNARead('id', query).reverseComplement().sequence
                 if rcSuffix:
-                    read.query_name += rcSuffix
+                    alignment.query_name += rcSuffix
 
             # Adjust the query id if it's a duplicate and we're not
             # allowing duplicates.
             if allowDuplicateIds:
-                queryId = read.query_name
+                queryId = alignment.query_name
             else:
-                count = idCount[read.query_name]
-                idCount[read.query_name] += 1
-                queryId = read.query_name + (
+                count = idCount[alignment.query_name]
+                idCount[alignment.query_name] += 1
+                queryId = alignment.query_name + (
                     '' if count == 0 else '/%d' % count)
 
-            referenceStart = read.reference_start
+            referenceStart = alignment.reference_start
             atStart = True
             queryIndex = 0
             referenceIndex = referenceStart
             alignedSequence = ''
 
-            for operation, length in read.cigartuples:
+            for operation, length in alignment.cigartuples:
 
                 # The operations are tested in the order they appear in
                 # https://samtools.github.io/hts-specs/SAMv1.pdf It would be
@@ -303,7 +308,11 @@ class PaddedSAM(object):
                 padChar * (referenceLength -
                            (referenceStart + len(alignedSequence))))
 
-            yield Read(queryId, paddedSequence)
+            read = Read(queryId, paddedSequence)
+            if alsoYieldAlignments:
+                yield (read, alignment)
+            else:
+                yield read
 
 
 @contextmanager


Test patch:
diff --git a/test/test_sam.py b/test/test_sam.py
index f8fd7e9..a6ee95a 100644
--- a/test/test_sam.py
+++ b/test/test_sam.py
@@ -706,3 +706,31 @@ class TestPaddedSAM(TestCase):
             queries = ps.queries()
             assertRaisesRegex(self, ValueError, error, list, queries)
             ps.close()
+
+    def testAlsoYieldAlignments(self):
+        """
+        A request for queries and their pysam alignments should have the
+        expected result.
+        """
+        data = '\n'.join([
+            '@SQ SN:ref1 LN:10',
+            'query1 0 ref1 2 60 2=2X2M * 0 0 TCTAGG 123456',
+            'query2 0 ref1 2 60 2= * 0 0 TC XY',
+        ]).replace(' ', '\t')
+
+        with dataFile(data) as filename:
+            ps = PaddedSAM(filename)
+            ((read1, alignment1),
+             (read2, alignment2)) = list(ps.queries(alsoYieldAlignments=True))
+
+            self.assertEqual(Read('query1', '-TCTAGG---'), read1)
+            self.assertEqual('TCTAGG', alignment1.query_sequence)
+            self.assertEqual('123456', ''.join(
+                map(lambda x: chr(x + 33), alignment1.query_qualities)))
+
+            self.assertEqual(Read('query2', '-TC-------'), read2)
+            self.assertEqual('TC', alignment2.query_sequence)
+            self.assertEqual('XY', ''.join(
+                map(lambda x: chr(x + 33), alignment2.query_qualities)))
+
+            ps.close()