Repository: cogent3/cogent3
Pull request #1887 (2024-06-05 07:04:34)

ENH: new_moltype.Moltype supports is_gapped and is_degenerate


Patch:
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index d3f781a9d..d4a9ed47f 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -480,6 +480,50 @@ class MolType:
         """reverse reverse complement of a sequence"""
         return self.complement(seq)[::-1]
 
+    @functools.singledispatchmethod
+    def is_degenerate(self, seq) -> bool:
+        """checks if a sequence contains degenerate characters"""
+        raise TypeError(f"{type(seq)} not supported")
+
+    @is_degenerate.register
+    def _(self, seq: bytes) -> bool:
+        return self.is_degenerate(seq.decode("utf8"))
+
+    @is_degenerate.register
+    def _(self, seq: str) -> bool:
+        return any(c in self.ambiguities for c in seq)
+
+    @is_degenerate.register
+    def _(self, seq: numpy.ndarray) -> bool:
+        # what index is the first degenerate character
+        for index, val in enumerate(self.degen_gapped_alphabet):
+            if val in self.ambiguities:
+                break
+        else:
+            return False
+        return (seq >= index).any()
+
+    @functools.singledispatchmethod
+    def is_gapped(self, seq) -> bool:
+        """checks if a sequence contains gaps"""
+        raise TypeError(f"{type(seq)} not supported")
+
+    @is_gapped.register
+    def _(self, seq: str) -> bool:
+        return self.gapped_alphabet.gap_char in seq
+
+    @is_gapped.register
+    def _(self, seq: bytes) -> bool:
+        return self.is_gapped(seq.decode("utf8"))
+
+    @is_gapped.register
+    def _(self, seq: bytes) -> bool:
+        return self.is_gapped(seq.decode("utf8"))
+
+    @is_gapped.register
+    def _(self, seq: numpy.ndarray) -> bool:
+        return (seq == self.degen_gapped_alphabet.gap_index).any()
+
 
 def _make_moltype_dict() -> dict[str, MolType]:
     """make a dictionary of local name space molecular types"""


Test patch:
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index 0e1c08a77..09d8d60ac 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -1,3 +1,4 @@
+import numpy
 import pytest
 
 from cogent3.core import moltype, new_moltype
@@ -77,3 +78,82 @@ def test_complement(name, seq):
     expect = "TGGGC" if name == "dna" else "UGGGC"
     got = dna.complement(seq)
     assert got == expect
+
+
+def make_typed(seq, data_type, moltype):
+    if data_type is numpy.ndarray:
+        seq = moltype.degen_gapped_alphabet.to_indices(seq)
+    elif data_type is bytes:
+        seq = seq.encode("utf-8")
+    return seq
+
+
+@pytest.mark.parametrize("data_type", (str, bytes, numpy.ndarray))
+@pytest.mark.parametrize(
+    "seq",
+    (
+        "N",
+        "R",
+        "Y",
+        "GCAUGUAGCUCGUCAGUCAGUACGUGCASCUAG",
+        "ACGYAUGCUGYEWEWNFMNFUWBYBCWUYBCJWBEIWFUB",
+    ),
+)
+def test_is_degenerate(seq, data_type):
+    seq = make_typed(seq, data_type, new_moltype.RNA)
+    assert new_moltype.RNA.is_degenerate(seq)
+
+
+@pytest.mark.parametrize("data_type", (str, bytes, numpy.ndarray))
+@pytest.mark.parametrize(
+    "seq",
+    (
+        "",
+        "A",
+        "UACGCUACAUGUACGUCAGUGCUAGCUA",
+    ),
+)
+def test_is_not_degenerate(seq, data_type):
+    seq = make_typed(seq, data_type, new_moltype.RNA)
+    assert not new_moltype.RNA.is_degenerate(seq)
+
+
+def test_is_degenerate_invalid():
+    with pytest.raises(TypeError):
+        new_moltype.RNA.is_degenerate(list("GAG"))
+
+
+@pytest.mark.parametrize("data_type", (str, bytes, numpy.ndarray))
+@pytest.mark.parametrize(
+    "seq",
+    (
+        "-",
+        "Y-",
+        "GC--A",
+        "-ACGYA",
+    ),
+)
+def test_is_gapped(seq, data_type):
+    seq = make_typed(seq, data_type, new_moltype.RNA)
+    assert new_moltype.RNA.is_gapped(seq)
+
+
+@pytest.mark.parametrize("data_type", (str, bytes, numpy.ndarray))
+@pytest.mark.parametrize(
+    "seq",
+    (
+        "",
+        "Y",
+        "GCA",
+        "ACGYA",
+    ),
+)
+def test_not_is_gapped(seq, data_type):
+    seq = make_typed(seq, data_type, new_moltype.RNA)
+    assert not new_moltype.RNA.is_gapped(seq)
+
+
+@pytest.mark.parametrize("moltype", (new_moltype.DNA, new_moltype.RNA))
+def test_gap_index_constant(moltype):
+    # make sure gap index is always the same
+    assert moltype.gapped_alphabet.gap_index == moltype.degen_gapped_alphabet.gap_index