Repository: cogent3/cogent3
Pull request #1891 (2024-06-07 01:35:35)

Enhancing alphabet and moltype


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 063d9d645..c5dddc219 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -9,6 +9,7 @@ import numpy
 
 
 StrORBytes = typing.Union[str, bytes]
+StrORBytesORArray = typing.Union[str, bytes, numpy.ndarray]
 OptInt = typing.Optional[int]
 OptStr = typing.Optional[str]
 OptBytes = typing.Optional[bytes]
@@ -114,7 +115,7 @@ class MonomerAlphabetABC(ABC):
     def from_indices(self, seq: numpy.ndarray) -> str: ...
 
     @abstractmethod
-    def to_bytes(self) -> bytes: ...
+    def as_bytes(self) -> bytes: ...
 
 
 def get_array_type(num_elements: int):
@@ -193,7 +194,7 @@ class array_to_bytes:
 
     def __call__(self, seq: numpy.ndarray) -> bytes:
         b = self._converter(seq.tobytes())
-        return bytearray(b)
+        return bytes(bytearray(b))
 
 
 class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
@@ -212,7 +213,7 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         self._gap_index = self.index(gap) if gap else None
         self.dtype = get_array_type(len(self))
         self._chars = set(self)  # for quick lookup
-        byte_chars = self.to_bytes()
+        byte_chars = self.as_bytes()
         self._bytes2arr = bytes_to_array(byte_chars, dtype=self.dtype)
         self._arr2bytes = array_to_bytes(byte_chars)
 
@@ -230,7 +231,7 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         return 1
 
     @functools.singledispatchmethod
-    def to_indices(self, seq: str) -> numpy.ndarray:
+    def to_indices(self, seq: StrORBytesORArray) -> numpy.ndarray:
         raise TypeError(f"{type(seq)} is invalid")
 
     @to_indices.register
@@ -239,7 +240,7 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
 
     @to_indices.register
     def _(self, seq: str) -> numpy.ndarray:
-        return self._bytes2arr(seq.encode("utf8"))
+        return self.to_indices(seq.encode("utf8"))
 
     @to_indices.register
     def _(self, seq: numpy.ndarray) -> numpy.ndarray:
@@ -282,23 +283,32 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         return KmerAlphabet(words=words, monomers=self, gap=gap, k=k)
 
     @functools.singledispatchmethod
-    def is_valid(self, seq) -> bool:
+    def is_valid(self, seq: StrORBytesORArray) -> bool:
         raise TypeError(f"{type(seq)} is invalid")
 
     @is_valid.register
     def _(self, seq: str) -> bool:
-        return set(seq) <= self._chars
+        return self.is_valid(self.to_indices(seq))
+
+    @is_valid.register
+    def _(self, seq: bytes) -> bool:
+        return self.is_valid(self.to_indices(seq))
 
     @is_valid.register
     def _(self, seq: numpy.ndarray) -> bool:
         return seq.min() >= 0 and seq.max() < len(self)
 
-    def to_bytes(self) -> bytes:
+    def as_bytes(self) -> bytes:
+        """returns self as a byte string"""
         if not isinstance(self[0], str):
             return bytes(bytearray(self))
 
         return "".join(self).encode("utf8")
 
+    def array_to_bytes(self, seq: numpy.ndarray) -> bytes:
+        """returns seq as a byte string"""
+        return self._arr2bytes(seq)
+
 
 @numba.jit(nopython=True)
 def coord_conversion_coeffs(num_states, k, dtype=None):  # pragma: no cover
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index baac710c2..4950f9bc9 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -15,6 +15,7 @@ OptStr = typing.Optional[str]
 OptCallable = typing.Optional[typing.Callable]
 SeqStrType = typing.Union[list[str], tuple[str, ...]]
 StrORBytes = typing.Union[str, bytes]
+StrORBytesORArray = typing.Union[str, bytes, numpy.ndarray]
 SeqStrBytesType = typing.Union[list[StrORBytes], tuple[StrORBytes, ...]]
 StrORArray = typing.Union[str, numpy.ndarray]
 
@@ -378,7 +379,7 @@ class MolType:
             # assume we have a nucleic acid moltype
             dest = "".join(complements[c] for c in self.degen_gapped_alphabet)
             self._complement = new_alphabet.convert_alphabet(
-                self.degen_gapped_alphabet.to_bytes(),
+                self.degen_gapped_alphabet.as_bytes(),
                 dest.encode("utf8"),
             )
 
@@ -463,7 +464,7 @@ class MolType:
         return self._make_seq(moltype=self, seq=seq or "", name=name, **kwargs)
 
     @functools.singledispatchmethod
-    def complement(self, seq) -> str:
+    def complement(self, seq: StrORBytesORArray) -> str:
         """converts a string or bytes into it's nucleic acid complement"""
         raise TypeError(f"{type(seq)} not supported")
 
@@ -475,22 +476,26 @@ class MolType:
     def _(self, seq: bytes) -> str:
         return self._complement(seq).decode("utf8")
 
+    @complement.register
+    def _(self, seq: numpy.ndarray) -> str:
+        return self.complement(self.degen_gapped_alphabet.array_to_bytes(seq))
+
     def rc(self, seq: str) -> str:
         """reverse reverse complement of a sequence"""
         return self.complement(seq)[::-1]
 
     @functools.singledispatchmethod
-    def is_degenerate(self, seq) -> bool:
+    def is_degenerate(self, seq: StrORBytesORArray) -> 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"))
+        return self.is_degenerate(self.degen_gapped_alphabet.to_indices(seq))
 
     @is_degenerate.register
     def _(self, seq: str) -> bool:
-        return any(c in self.ambiguities for c in seq)
+        return self.is_degenerate(self.degen_gapped_alphabet.to_indices(seq))
 
     @is_degenerate.register
     def _(self, seq: numpy.ndarray) -> bool:


Test patch:
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index 86721902d..165980920 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -41,7 +41,7 @@ def test_convert_alphabet_mixed_length():
 )
 def test_to_bytes(moltype):
     alpha = moltype.alphabet
-    got = alpha.to_bytes()
+    got = alpha.as_bytes()
     expect = "".join(alpha).encode("utf8")
     assert len(got) == len(alpha)
     assert got == expect
@@ -49,7 +49,7 @@ def test_to_bytes(moltype):
 
 def test_to_bytes_bytes():
     alpha = new_moltype.BYTES.alphabet
-    got = alpha.to_bytes()
+    got = alpha.as_bytes()
     assert len(got) == len(alpha)
     expect = bytes(bytearray(range(256)))
     assert got == expect
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index 09d8d60ac..368c1d05b 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -71,7 +71,9 @@ def test_str_moltype():
     assert text == f"MolType({tuple('TCAG')})"
 
 
-@pytest.mark.parametrize("seq", ("ACCCG", b"ACCCG"))
+@pytest.mark.parametrize(
+    "seq", ("ACCCG", b"ACCCG", numpy.array([2, 1, 1, 1, 3], dtype=numpy.uint8))[-1:]
+)
 @pytest.mark.parametrize("name", ("dna", "rna"))
 def test_complement(name, seq):
     dna = new_moltype.get_moltype(name)
@@ -95,6 +97,7 @@ def make_typed(seq, data_type, moltype):
         "N",
         "R",
         "Y",
+        "?",  # IUPAC missing is also a degenerate
         "GCAUGUAGCUCGUCAGUCAGUACGUGCASCUAG",
         "ACGYAUGCUGYEWEWNFMNFUWBYBCWUYBCJWBEIWFUB",
     ),