Merge pull request #1092 from ajcr/add_allele_variant_types

Add VariantRecord.alleles_variant_types
pysam-developers · May 26, 2022 · a7440c7 · a7440c7
2 parents 5b80eb4 + 18c3528
commit a7440c7
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 0 deletions.
diff --git a/pysam/libcbcf.pyx b/pysam/libcbcf.pyx
@@ -3261,6 +3261,7 @@ cdef class VariantRecord(object):
             self.ptr.rlen = rlen
         else:
             self.ptr.rlen = len(values[0])
+        r.d.var_type = -1
         bcf_sync_end(self)
 
     @property
@@ -3289,6 +3290,7 @@ cdef class VariantRecord(object):
             raise ValueError('cannot set null alt allele')
         ref  = [r.d.allele[0] if r.d.allele and r.n_allele else b'.']
         self.alleles = ref + value
+        r.d.var_type = -1
 
     @property
     def filter(self):
@@ -3318,6 +3320,34 @@ cdef class VariantRecord(object):
             raise ValueError('Error unpacking VariantRecord')
         return makeVariantRecordSamples(self)
 
+    property alleles_variant_types:
+        def __get__(self):
+            cdef bcf1_t *r = self.ptr
+            cdef tuple result = PyTuple_New(r.n_allele)
+
+            for i in range(r.n_allele):
+                tp = bcf_get_variant_type(r, i)
+
+                if tp == VCF_REF:
+                    v_type = "REF"
+                elif tp == VCF_SNP:
+                    v_type = "SNP"
+                elif tp == VCF_MNP:
+                    v_type = "MNP"
+                elif tp == VCF_INDEL:
+                    v_type = "INDEL"
+                elif tp == VCF_BND:
+                    v_type = "BND"
+                elif tp == VCF_OVERLAP:
+                    v_type = "OVERLAP"
+                else:
+                    v_type = "OTHER"
+
+                PyTuple_SET_ITEM(result, i, v_type)
+                Py_INCREF(v_type)
+
+            return result
+
     def __richcmp__(VariantRecord self not None, VariantRecord other not None, int op):
         if op != 2 and op != 3:
             return NotImplemented

diff --git a/pysam/libchtslib.pxd b/pysam/libchtslib.pxd
@@ -1558,6 +1558,9 @@ cdef extern from "htslib/vcf.h" nogil:
     uint8_t VCF_MNP
     uint8_t VCF_INDEL
     uint8_t VCF_OTHER
+    uint8_t VCF_BND
+    uint8_t VCF_OVERLAP
+
 
     ctypedef struct variant_t:
         int type, n     # variant type and the number of bases affected, negative for deletions

diff --git a/tests/VariantFile_test.py b/tests/VariantFile_test.py
@@ -364,6 +364,23 @@ def testAlleles(self):
         self.assertEqual(alleles, [
                          ('T',), ('G', 'A'), ('T', 'A'), ('A', 'G', 'T'), ('GTCT', 'G', 'GTACT')])
 
+    def testAllelesVariantTypes(self):
+        fn = os.path.join(CBCF_DATADIR, self.filename)
+        v = pysam.VariantFile(fn)
+        rec = next(v)
+
+        self.assertEqual(rec.alleles, ('T',))
+        self.assertEqual(rec.alleles_variant_types, ("REF",))
+
+        rec.alleles = ("T", "C")
+        self.assertEqual(rec.alleles_variant_types, ("REF", "SNP"))
+
+        rec.alts = ("TC",)
+        self.assertEqual(rec.alleles_variant_types, ("REF", "INDEL"))
+
+        rec.ref = "AG"
+        self.assertEqual(rec.alleles_variant_types, ("REF", "MNP"))
+
     def testQual(self):
         fn = os.path.join(CBCF_DATADIR, self.filename)
         v = pysam.VariantFile(fn)