Recode getSequenceInRange() to return str directly

Also optimise it to unpack high and low nibbles explicitly.

Author: jmarshall

pysam/libcalignedsegment.pyx

@@ -65,7 +65,7 @@ from cpython cimport array as c_array
 from cpython.bytearray cimport PyByteArray_FromStringAndSize, PyByteArray_GET_SIZE
 from cpython.bytes cimport PyBytes_FromStringAndSize
 from cpython.unicode cimport PyUnicode_DATA, PyUnicode_1BYTE_DATA, PyUnicode_1BYTE_KIND, \
-    PyUnicode_GET_LENGTH, PyUnicode_KIND, PyUnicode_WRITE
+    PyUnicode_GET_LENGTH, PyUnicode_KIND, PyUnicode_New, PyUnicode_WRITE
 from libc.string cimport memset, strchr
 from libc.stdint cimport INT8_MIN, INT16_MIN, INT32_MIN, \
     INT8_MAX, INT16_MAX, INT32_MAX, \
@@ -632,29 +632,31 @@ cdef inline int32_t getQueryEnd(bam1_t *src) except -1:
     return end_offset
 
 
-cdef inline bytes getSequenceInRange(bam1_t *src,
-                                     uint32_t start,
-                                     uint32_t end):
+cdef inline str getSequenceInRange(const bam1_t *src, uint32_t start, uint32_t end):
     """return python string of the sequence in a bam1_t object.
     """
+    seq = PyUnicode_New(end - start, 127)
+    cdef char *dest = <char *> PyUnicode_1BYTE_DATA(seq)
+    cdef unsigned int di = 0
 
-    cdef uint8_t * p
-    cdef uint32_t k
-    cdef char * s
+    cdef const uint8_t *p = bam_get_seq(src)
+    cdef unsigned int i = start // 2
 
-    if not src.core.l_qseq:
-        return None
+    if start & 1:
+        dest[di] = seq_nt16_str[p[i] & 0x0f]
+        i += 1
+        di += 1
 
-    seq = PyBytes_FromStringAndSize(NULL, end - start)
-    s   = <char*>seq
-    p   = pysam_bam_get_seq(src)
+    for _ in range(i, end // 2):
+        dest[di]   = seq_nt16_str[p[i] >> 4]
+        dest[di+1] = seq_nt16_str[p[i] & 0x0f]
+        i += 1
+        di += 2
 
-    for k from start <= k < end:
-        # equivalent to seq_nt16_str[bam1_seqi(s, i)] (see bam.c)
-        # note: do not use string literal as it will be a python string
-        s[k-start] = seq_nt16_str[p[k//2] >> 4 * (1 - k%2) & 0xf]
+    if end & 1:
+        dest[di] = seq_nt16_str[p[i] >> 4]
 
-    return charptr_to_bytes(seq)
+    return seq
 
 
 #####################################################################
@@ -1530,8 +1532,7 @@ cdef class AlignedSegment:
                 self.cache.query_sequence = None
                 return None
 
-            self.cache.query_sequence = force_str(getSequenceInRange(
-                src, 0, src.core.l_qseq))
+            self.cache.query_sequence = getSequenceInRange(src, 0, src.core.l_qseq)
             return self.cache.query_sequence
 
         def __set__(self, seq):
@@ -1898,7 +1899,7 @@ cdef class AlignedSegment:
             start = getQueryStart(src)
             end   = getQueryEnd(src)
 
-            self.cache.query_alignment_sequence = force_str(getSequenceInRange(src, start, end))
+            self.cache.query_alignment_sequence = getSequenceInRange(src, start, end)
             return self.cache.query_alignment_sequence
 
     property query_alignment_qualities:

tests/AlignedSegment_bench.py

@@ -35,3 +35,14 @@ def test_reverse_complement(pytestconfig, benchmark):
     if pytestconfig.get_verbosity() > 0: print(f"LEN={5 * count} ", end="", flush=True)
     result = benchmark(pysam.reverse_complement, "AATGC" * count)
     assert result == "GCATT" * count
+
+
+def test_query_sequence(pytestconfig, benchmark):
+    seq = "AATGC" * (int(os.environ.get("LEN", 300)) // 5)
+    if pytestconfig.get_verbosity() > 0: print(f"LEN={len(seq)} ", end="", flush=True)
+
+    read = pysam.AlignedSegment()
+    read.query_sequence = seq
+    # What we actually need for meaningful benchmarking is a non-caching non-property to call
+    result = benchmark(lambda : read.query_sequence)
+    assert result == seq

tests/AlignedSegment_test.py

@@ -1951,6 +1951,35 @@ def test_longarray_to_qualstr(self):
             pysam.array_to_qualitystring(qual_array)
 
 
+@pytest.mark.parametrize("seq", [
+    "A",
+    "AT",
+    "GCA",
+    "ATGC",
+    "AATTG",
+    pytest.param("ABCDGHKMNRSTVWY", id="iupac"),
+])
+def test_sequence_unpacking(seq):
+    a = pysam.AlignedSegment()
+    a.query_sequence = seq
+    assert a.query_sequence == seq
+
+
+@pytest.mark.parametrize("start,stop", [
+    (0, 0), (1, 0), (2, 0), (3, 0),
+    (0, 1), (1, 1), (2, 1), (3, 1),
+    (0, 2), (1, 2), (2, 2), (3, 2),
+    (0, 3), (1, 3), (2, 3), (3, 3),
+])
+def test_subsequence_unpacking(start, stop):
+    a = pysam.AlignedSegment()
+    a.query_sequence = "ABCDGHKMNRSTVWY"
+    start_clip = f"{start}S" if start > 0 else ""
+    stop_clip = f"{stop}S" if stop > 0 else ""
+    a.cigarstring = f"{start_clip}{a.query_length - start - stop}M{stop_clip}"
+    assert a.query_alignment_sequence == a.query_sequence[start : a.query_length - stop]
+
+
 @pytest.mark.parametrize("seq,revcomp", [
     ("", ""),
     ("A", "T"),