Repository: cogent3/cogent3
Pull request #1889 (2024-06-06 03:21:03)

Monomer alphabets to_bytes() method simplifies building alphabet inter converters


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index d6394d206..063d9d645 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -113,6 +113,9 @@ class MonomerAlphabetABC(ABC):
     @abstractmethod
     def from_indices(self, seq: numpy.ndarray) -> str: ...
 
+    @abstractmethod
+    def to_bytes(self) -> bytes: ...
+
 
 def get_array_type(num_elements: int):
     """Returns the smallest numpy integer dtype that can contain elements
@@ -209,10 +212,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
-        try:
-            byte_chars = b"".join([c.encode("utf8") for c in chars])
-        except AttributeError:
-            byte_chars = bytes(bytearray(chars))
+        byte_chars = self.to_bytes()
         self._bytes2arr = bytes_to_array(byte_chars, dtype=self.dtype)
         self._arr2bytes = array_to_bytes(byte_chars)
 
@@ -293,6 +293,12 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
     def _(self, seq: numpy.ndarray) -> bool:
         return seq.min() >= 0 and seq.max() < len(self)
 
+    def to_bytes(self) -> bytes:
+        if not isinstance(self[0], str):
+            return bytes(bytearray(self))
+
+        return "".join(self).encode("utf8")
+
 
 @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 d4a9ed47f..81ffbe6c3 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -376,10 +376,9 @@ class MolType:
         )
         if complements:
             # assume we have a nucleic acid moltype
-            src = "".join(self._degen_gapped)
-            dest = "".join(complements[c] for c in self._degen_gapped)
+            dest = "".join(complements[c] for c in self.degen_gapped_alphabet)
             self._complement = new_alphabet.convert_alphabet(
-                src.encode("utf8"),
+                self.degen_gapped_alphabet.to_bytes(),
                 dest.encode("utf8"),
             )
 


Test patch:
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index 5992b76f4..86721902d 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -36,6 +36,26 @@ def test_convert_alphabet_mixed_length():
         new_alphabet.convert_alphabet(src="ACG", dest="TGCA")
 
 
+@pytest.mark.parametrize(
+    "moltype", (new_moltype.DNA, new_moltype.DNA, new_moltype.ASCII)
+)
+def test_to_bytes(moltype):
+    alpha = moltype.alphabet
+    got = alpha.to_bytes()
+    expect = "".join(alpha).encode("utf8")
+    assert len(got) == len(alpha)
+    assert got == expect
+
+
+def test_to_bytes_bytes():
+    alpha = new_moltype.BYTES.alphabet
+    got = alpha.to_bytes()
+    assert len(got) == len(alpha)
+    expect = bytes(bytearray(range(256)))
+    assert got == expect
+    assert chr(got[65]) == chr(65)
+
+
 @pytest.mark.parametrize("words", ("AACG", ["A", "CG"], ["CC", "AA"], "", (), []))
 def test_invalid_words(words):
     with pytest.raises(ValueError):