fix: allow for IUPAC codes in GetCanonicalMotif()

trtools/utils/tests/test_utils.py

@@ -115,6 +115,8 @@ def test_GetCanonicalMotif():
     assert(utils.GetCanonicalMotif("TTGTT")=="AAAAC")
     assert(utils.GetCanonicalMotif("")=="")
     assert(utils.GetCanonicalMotif("cag")=="AGC")
+    assert(utils.GetCanonicalMotif("AARRG")=="AARRG")
+    assert(utils.GetCanonicalMotif("YARRG")=="ARRGY")
 
 # GetCanonicalOneStrand
 def test_GetCanonicalOneStrand():
@@ -126,13 +128,25 @@ def test_GetCanonicalOneStrand():
     assert(utils.GetCanonicalOneStrand("TTGTT")=="GTTTT")
     assert(utils.GetCanonicalOneStrand("")=="")
     assert(utils.GetCanonicalOneStrand("at")=="AT")
+    # Additional tests with IUPAC codes
+    assert(utils.GetCanonicalOneStrand("RY")=="RY")
+    assert(utils.GetCanonicalOneStrand("YR")=="RY")
+    assert(utils.GetCanonicalOneStrand("SW")=="SW")
+    assert(utils.GetCanonicalOneStrand("WS")=="SW")
+    assert(utils.GetCanonicalOneStrand("KM")=="KM")
+    assert(utils.GetCanonicalOneStrand("MK")=="KM")
 
 # ReverseComplement
 def test_ReverseComplement():
     assert(utils.ReverseComplement("CGAT")=="ATCG")
     assert(utils.ReverseComplement("")=="")
     assert(utils.ReverseComplement("CGNT")=="ANCG")
     assert(utils.ReverseComplement("ccga")=="TCGG")
+    # additional tests with IUPAC codes
+    assert(utils.ReverseComplement("RYASWKM")=="KMWSTRY")
+    # also test the characters that don't change
+    assert(utils.ReverseComplement("BDHV")=="BDHV")
+    assert(utils.ReverseComplement("N")=="N")
 
 # InferRepeatSequence
 def test_InferRepeatSequence():

trtools/utils/utils.py

@@ -14,7 +14,6 @@
 
 import trtools.utils.common as common # pragma: no cover
 
-nucToNumber={"A":0,"C":1,"G":2,"T":3}
 
 def LoadSingleReader(
         vcf_loc: str,
@@ -387,9 +386,9 @@ def GetCanonicalMotif(repseq):
     repseq_r = GetCanonicalOneStrand(ReverseComplement(repseq))
     # choose first seq alphabetically
     for i in range(len(repseq_f)):
-        if nucToNumber[repseq_f[i]] < nucToNumber[repseq_r[i]]:
+        if repseq_f[i] < repseq_r[i]:
             return repseq_f
-        if nucToNumber[repseq_r[i]] < nucToNumber[repseq_f[i]]:
+        if repseq_r[i] < repseq_f[i]:
             return repseq_r
     return repseq_f
 
@@ -420,16 +419,16 @@ def GetCanonicalOneStrand(repseq):
     for i in range(size):
         newseq = repseq[size-i:]+repseq[0:size-i]
         for j in range(size):
-            if nucToNumber[newseq[j]] < nucToNumber[canonical[j]]:
+            if newseq[j] < canonical[j]:
                 canonical = newseq
-            elif nucToNumber[newseq[j]] > nucToNumber[canonical[j]]:
+            elif newseq[j] > canonical[j]:
                 break
     return canonical
 
 def ReverseComplement(seq):
     r"""Get reverse complement of a sequence.
 
-    Converts everything to uppsercase.
+    Converts everything to uppercase and handles IUPAC codes.
 
     Parameters
     ----------
@@ -445,21 +444,19 @@ def ReverseComplement(seq):
     --------
     >>> ReverseComplement("AGGCT")
     'AGCCT'
+    >>> ReverseComplement("AGGCTRY")
+    'RYAGCCT'
     """
+    iupac_complement = {
+        'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A',
+        'R': 'Y', 'Y': 'R', 'S': 'S', 'W': 'W',
+        'K': 'M', 'M': 'K', 'B': 'V', 'D': 'H',
+        'H': 'D', 'V': 'B', 'N': 'N'
+    }
     seq = seq.upper()
     newseq = ""
-    size = len(seq)
-    for i in range(len(seq)):
-        char = seq[len(seq)-i-1]
-        if char == "A":
-            newseq += "T"
-        elif char == "G":
-            newseq += "C"
-        elif char == "C":
-            newseq += "G"
-        elif char == "T":
-            newseq += "A"
-        else: newseq += "N"
+    for char in reversed(seq):
+        newseq += iupac_complement.get(char, 'N')
     return newseq
 
 def InferRepeatSequence(seq, period):