Fix corner case bug for the sequence buffer iterator

It was possible that the the remaining sequence buffer exactly aligned
with the end of the sequence. This cause some sequence to be double
counted.

Author: EricR86

newmap/fasta.py

@@ -38,51 +38,157 @@ def sequence_segments(
     # NB: Mutable sequence of bytes
     # NB: This is over a 1000x (not a typo) speed-up over a byte object
     working_sequence_buffer = bytearray()
+    overlap_buffer = bytearray()
 
-    # NB: Line buffered reading is probably the best way to handle edge cases
-    # around sequence ID parsing and newline characters
-    # NB: fasta_file is assumed to be opened in binary mode (rb)
-    for fasta_line in fasta_file:
-        # If we are on a new sequence (sequence ID)
-        # NB: Assume that either of the delimiters are indicators of a
-        # new sequence, notably including comments
-        if fasta_line.startswith(FASTA_FILE_IGNORE_DELIMITERS):  # type: ignore
-            # Yield the current sequence segment if there is remaining sequence
-            # NB: We always keep the lookahead/overlap in the working sequence
-            # buffer, therefore there can only be sequence remaining if it is
-            # longer than the lookahead/overlap length
-            if len(working_sequence_buffer) > sequence_overlap_length:
-                yield SequenceSegment(current_sequence_id,  # type: ignore
-                                      bytes(working_sequence_buffer))
-
-            # Get the new reference sequence name
+    sequences = []  # working list of sequences
+
+    # import debugpy
+    # debugpy.listen(5678)
+    # debugpy.wait_for_client()
+    # debugpy.breakpoint()
+    # For every line in the fasta file
+    for line in fasta_file:
+        line = line.rstrip()  # Remove trailing newline
+
+        # While there is enough working sequence buffer to fill the requested
+        # sequence length
+        # Create sequences for each segment
+        sequences.extend(get_sequences_from_buffer(
+                             working_sequence_buffer,
+                             overlap_buffer,
+                             sequence_length,
+                             sequence_overlap_length))
+
+        # If the current line is a sequence ID
+        if line.startswith(FASTA_FILE_IGNORE_DELIMITERS):  # type: ignore
+            # Yield the remaining sequences
+            for sequence_buffer in get_remaining_sequence_segments(
+                                   current_sequence_id,  # type: ignore
+                                   sequences,
+                                   working_sequence_buffer,
+                                   overlap_buffer,
+                                   sequence_length,
+                                   sequence_overlap_length):
+                yield sequence_buffer
+
+            # Empty working sequences
+            sequences.clear()
+
+            # Update the working sequence ID
             # NB: remove leading '>'
-            current_sequence_id = fasta_line.split()[0][1:]  # type: ignore
-            # Reset the working sequence buffer
-            working_sequence_buffer = bytearray()
+            current_sequence_id = line.split()[0][1:]  # type: ignore
+            # Clear the overlap buffer
+            overlap_buffer.clear()
+            # Clear the working buffer
+            working_sequence_buffer.clear()
 
-        # Otherwise the line we are on is sequence data
+        # Otherwise the line is not a sequence ID and is sequence data
         else:
-            fasta_line = fasta_line.rstrip()  # Remove trailing newline
-            # Add to the working sequence buffer
-            working_sequence_buffer += fasta_line  # type: ignore
-            # While we have enough sequence buffer to fill a sequence segment
-            while len(working_sequence_buffer) >= sequence_length:
-                yield SequenceSegment(
-                    current_sequence_id,  # type: ignore
-                    bytes(working_sequence_buffer[:sequence_length]))
-                # Truncate the working sequence buffer by the sequence length
-                # minus the lookahead
-                # XXX: Assert that the kmer/sequence length is always larger
-                # than the lookahead length?
-                truncate_length = sequence_length - sequence_overlap_length
-                working_sequence_buffer = \
-                    working_sequence_buffer[truncate_length:]
-
-    # Yield the last sequence segment
-    # NB: We always keep the lookahead in the working sequence buffer
-    # So there needs to be check if it is longer the lookahead length
-    if len(working_sequence_buffer) > sequence_overlap_length:
-        yield SequenceSegment(current_sequence_id,  # type: ignore
-                              bytes(working_sequence_buffer),
-                              epilogue=True)
+            # If any sequences were created
+            if sequences:
+                # Create all sequence segments but for the last
+                for sequence in sequences[:-1]:
+                    # Yield a sequence segment without the epilogue flag set
+                    yield SequenceSegment(current_sequence_id, bytes(sequence))
+                # Carry over the last sequence to the next iteration
+                # in case this the last line it the sequence filled the
+                # remaining buffer exactly
+                sequences = [sequences[-1]]
+
+            # Add the sequence line to the working buffer
+            working_sequence_buffer += line  # type: ignore
+
+    # Yield the remaining sequences
+    for sequence_buffer in get_remaining_sequence_segments(
+                           current_sequence_id,  # type: ignore
+                           sequences,
+                           working_sequence_buffer,
+                           overlap_buffer,
+                           sequence_length,
+                           sequence_overlap_length):
+
+        yield sequence_buffer
+
+
+def get_sequences_from_buffer(working_sequence_buffer: bytearray,
+                              overlap_sequence: bytearray,
+                              sequence_length: int,
+                              sequence_overlap_length: int):
+    """Returns a list of overlapping byte sequences from a sequence buffer.
+       Modifies the working sequence buffer and overlap buffer in place.
+    """
+
+    # If there is no sequence buffer
+    if not working_sequence_buffer:
+        # Return nothing
+        return []
+
+    sequences = []
+
+    non_overlap_length = sequence_length - sequence_overlap_length
+
+    while (len(working_sequence_buffer) + len(overlap_sequence) >=
+           sequence_length):
+        # If there is an overlap buffer
+        if overlap_sequence:
+            # Create the sequence with the overlap
+            sequence = bytes(
+                overlap_sequence +
+                working_sequence_buffer[:non_overlap_length])
+            bytes_used = non_overlap_length
+        else:
+            # Otherwise create the sequence without the overlap
+            sequence = bytes(working_sequence_buffer[:sequence_length])
+            bytes_used = sequence_length
+
+        # Add to our working list of sequences
+        sequences.append(sequence)
+        # Update the overlap buffer if it exists by taking the last
+        # current calculated sequence
+        if sequence_overlap_length:
+            # NB: Avoid re-assignment to modifiy in place
+            overlap_sequence[:] = sequence[-sequence_overlap_length:]
+        # Truncate the start of working sequence buffer by bytes used
+        working_sequence_buffer[:bytes_used] = b''
+
+    return sequences
+
+
+def get_remaining_sequence_segments(sequence_id: bytes,
+                                    sequences: list[bytes],
+                                    working_sequence_buffer: bytearray,
+                                    overlap_buffer: bytearray,
+                                    sequence_length: int,
+                                    sequence_overlap_length: int):
+    # Assumes last of any buffer is the the epilogue
+
+    sequence_segments = []
+
+    sequences.extend(get_sequences_from_buffer(
+                         working_sequence_buffer,
+                         overlap_buffer,
+                         sequence_length,
+                         sequence_overlap_length))
+
+    if working_sequence_buffer:
+        sequences.append(bytes(overlap_buffer + working_sequence_buffer))
+
+    # If any sequences were created
+    if sequences:
+        # Create a sequence segment for all but the last element
+        # NB: Empty on a single list
+        for sequence in sequences[:-1]:
+            sequence_segments.append(
+                    SequenceSegment(sequence_id, bytes(sequence))
+            )
+
+        # Create a sequence segment for the last element with the
+        # epilogue flag set
+
+        sequence_segments.append(
+            SequenceSegment(sequence_id,
+                            bytes(sequences[-1]),
+                            epilogue=True)
+        )
+
+    return sequence_segments

tests/test_sequence_buffer_iter.py

@@ -20,6 +20,7 @@ def test_entire_sequence(self):
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'CGCANCAGAGCANCGNCG', sequence_buffer.data)
+        self.assertTrue(sequence_buffer.epilogue)
         self.assertRaises(StopIteration, next, buffer_iter)
 
     def test_sequence_overlap(self):
@@ -28,26 +29,32 @@ def test_sequence_overlap(self):
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'CGCAN', sequence_buffer.data)
+        self.assertFalse(sequence_buffer.epilogue)
 
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'ANCAG', sequence_buffer.data)
+        self.assertFalse(sequence_buffer.epilogue)
 
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'AGAGC', sequence_buffer.data)
+        self.assertFalse(sequence_buffer.epilogue)
 
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'GCANC', sequence_buffer.data)
+        self.assertFalse(sequence_buffer.epilogue)
 
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'NCGNC', sequence_buffer.data)
+        self.assertFalse(sequence_buffer.epilogue)
 
         sequence_buffer = next(buffer_iter)
         self.assertEqual(b'chr2', sequence_buffer.id)
         self.assertEqual(b'NCG', sequence_buffer.data)
+        self.assertTrue(sequence_buffer.epilogue)
 
         self.assertRaises(StopIteration, next, buffer_iter)
 
@@ -160,3 +167,23 @@ def test_single_nucleotide_epilogue(self):
 
         self.assertRaises(StopIteration, next, buffer_iter)
         fasta_file.close()
+
+    def test_sizes_and_epilogues(self):
+        buffer_iter = sequence_segments(self.genome_fasta_file, 20, 0)
+
+        sequence_buffer = next(buffer_iter)
+        self.assertEqual(b'chr1', sequence_buffer.id)
+        self.assertEqual(len(sequence_buffer.data), 20)
+        self.assertTrue(sequence_buffer.epilogue)
+
+        sequence_buffer = next(buffer_iter)
+        self.assertEqual(b'chr2', sequence_buffer.id)
+        self.assertEqual(len(sequence_buffer.data), 20)
+        self.assertFalse(sequence_buffer.epilogue)
+
+        sequence_buffer = next(buffer_iter)
+        self.assertEqual(b'chr2', sequence_buffer.id)
+        self.assertEqual(len(sequence_buffer.data), 10)
+        self.assertTrue(sequence_buffer.epilogue)
+
+        self.assertRaises(StopIteration, next, buffer_iter)