Repository: cogent3/cogent3
Pull request #2048 (2024-10-17 00:16:46)

Update new CodonAlphabet and GeneticCodes
## Summary by Sourcery

Enhance the genetic code translation functionality by introducing an 'incomplete_ok' parameter to handle incomplete codons. Refactor the genetic code retrieval and error handling for improved clarity and robustness. Update tests to cover new functionality and ensure reliability of the genetic code and alphabet classes.

Enhancements:
- Add an 'incomplete_ok' parameter to the 'translate' function to handle incomplete codons more flexibly.
- Refactor the 'get_code' function to simplify the handling of 'code_id' parameter.
- Improve error messages by using more descriptive representations for invalid inputs.

Tests:
- Add tests for the new 'incomplete_ok' parameter in the 'translate' function to ensure correct handling of incomplete codons.
- Introduce tests for invalid codon indices to verify error handling.
- Add tests for the 'SenseCodonAlphabet' class to ensure correct serialization and deserialization.


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index c457a773a..3f9a5a6aa 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -19,10 +19,6 @@ OptStr = typing.Optional[str]
 OptBytes = typing.Optional[bytes]
 
 
-def _element_type(val: typing.Sequence) -> typing.Any:
-    return type(val[0])
-
-
 @functools.singledispatch
 def _coerce_to_type(orig: StrORBytes, text: str) -> StrORBytes:
     raise TypeError(f"{type(orig)} is invalid")
@@ -856,8 +852,8 @@ def deserialise_kmer_alphabet(data: dict) -> KmerAlphabet:
     return KmerAlphabet.from_rich_dict(data)
 
 
-class CodonAlphabet(tuple):
-    """represents the sense-codons of a GeneticCode"""
+class SenseCodonAlphabet(tuple):
+    """Alphabet for the sense-codons of a GeneticCode"""
 
     def __new__(
         cls,
@@ -905,7 +901,7 @@ class CodonAlphabet(tuple):
 
     @property
     def gap_index(self):
-        return self._to_indices(self.gap_char) if self._gap_char else None
+        return self._to_indices[self.gap_char] if self._gap_char else None
 
     @property
     def missing_char(self):
@@ -928,7 +924,15 @@ class CodonAlphabet(tuple):
 
     @to_indices.register
     def _(self, seq: list) -> numpy.ndarray:
-        return [self.to_index(c) for c in seq]
+        return numpy.array([self.to_index(c) for c in seq], dtype=self.dtype)
+
+    @to_indices.register
+    def _(self, seq: numpy.ndarray) -> numpy.ndarray:
+        # we assume that this is a dna sequence encoded as a numpy array
+        size = len(seq) // 3
+        return self.to_indices(
+            [self.monomers.from_indices(c) for c in seq.reshape(size, 3)]
+        )
 
     def to_index(self, codon: str) -> int:
         if len(codon) != 3:
@@ -936,7 +940,7 @@ class CodonAlphabet(tuple):
         try:
             return self._to_indices[codon]
         except KeyError as e:
-            raise ValueError(f"{codon=!r} not in alphabet") from e
+            raise AlphabetError(f"{codon=!r} not in alphabet") from e
 
     def from_index(self, index: int) -> str:
         if index > len(self) or index < 0:
@@ -964,7 +968,7 @@ class CodonAlphabet(tuple):
         try:
             _ = self.to_indices(seq)
             return True
-        except ValueError:
+        except (ValueError, AlphabetError):
             return False
 
     @is_valid.register
@@ -1000,9 +1004,9 @@ class CodonAlphabet(tuple):
         return cls(**data)
 
 
-@register_deserialiser(get_object_provenance(CodonAlphabet))
-def deserialise_codon_alphabet(data: dict) -> CodonAlphabet:
-    return CodonAlphabet.from_rich_dict(data)
+@register_deserialiser(get_object_provenance(SenseCodonAlphabet))
+def deserialise_codon_alphabet(data: dict) -> SenseCodonAlphabet:
+    return SenseCodonAlphabet.from_rich_dict(data)
 
 
 _alphabet_moltype_map = {}
diff --git a/src/cogent3/core/new_genetic_code.py b/src/cogent3/core/new_genetic_code.py
index b3253d307..041fb7b9e 100644
--- a/src/cogent3/core/new_genetic_code.py
+++ b/src/cogent3/core/new_genetic_code.py
@@ -12,6 +12,7 @@ versa, lists of codons that they produce will be provided in DNA format.
 import collections
 import contextlib
 import dataclasses
+import functools
 import itertools
 import typing
 
@@ -193,14 +194,18 @@ class GeneticCode:
             return self._aa_to_codon.get(item, set())
 
         if len(item) != 3:
-            raise InvalidCodonError(f"Codon or aa {item} has wrong length")
+            raise InvalidCodonError(f"Codon or aa {item!r} has wrong length")
 
         key = item.upper()
         key = key.replace("U", "T")
         return self._codon_to_aa.get(key, "X")
 
     def translate(
-        self, dna: StrORBytesORArray, start: int = 0, rc: bool = False
+        self,
+        dna: StrORBytesORArray,
+        start: int = 0,
+        rc: bool = False,
+        incomplete_ok: bool = True,
     ) -> str:
         """Translates DNA to protein.
 
@@ -212,17 +217,22 @@ class GeneticCode:
             position to begin translation (used to implement frames)
         rc
             if True, returns the translation of the reverse complement sequence
+        incomplete_ok
+            if True, translates codons that are a mix of gaps and bases
+            as a gap. If False, raises an AlphabetError on those incomplete
+            cases.
 
         Notes
         -----
         Sequences are truncated to be a multiple of 3.
         Codons containing ambiguous nucleotides are translated as 'X',
-        codons containing a gap character are translated as '-'. Codons with
-        a mix of ambiguous nucleotides and gaps are translated as 'X'.
+        codons containing a gap character are translated as '-' unless
+        incomplete_ok is False. Codons with a mix of ambiguous nucleotides
+        are translated as 'X'.
 
         Returns
         -------
-        String containing amino acid sequence.
+        The amino acid sequence as a string.
         """
         # refactor: design
         # previous implementation converted mix of ambigs and nt to '?', potentially
@@ -234,7 +244,11 @@ class GeneticCode:
             dna = dna[:-diff]
 
         # convert to indices and then bytes
-        seq = self.codons.to_indices(dna)
+        if incomplete_ok:
+            seq = self.codons.to_indices(dna)
+        else:
+            codons = self.get_alphabet(include_stop=True, include_gap=True)
+            seq = codons.to_indices(dna)
 
         if rc:
             return self._translate_minus(seq.tobytes()).decode("utf8")[::-1]
@@ -256,9 +270,10 @@ class GeneticCode:
         """Returns True if codon is a stop codon, False otherwise."""
         return self[codon] == "*"
 
+    @functools.cache
     def get_alphabet(
         self, include_gap: bool = False, include_stop: bool = False
-    ) -> new_alphabet.CodonAlphabet:
+    ) -> new_alphabet.SenseCodonAlphabet:
         """returns a codon alphabet
 
         Parameters
@@ -276,8 +291,8 @@ class GeneticCode:
         if include_gap:
             codons += (gap,)
 
-        return new_alphabet.CodonAlphabet(
-            words=codons, monomers=self.moltype.alphabet, gap=gap
+        return new_alphabet.SenseCodonAlphabet(
+            words=codons, monomers=self.moltype.degen_gapped_alphabet, gap=gap
         )
 
 
@@ -471,14 +486,13 @@ def get_code(code_id: StrORInt = 1) -> GeneticCode:
     # refactor: simplify
     # added for compatibility with previous API, should be removed when
     # support for old style is dropped
-    if code_id is None:
-        code_id = 1
+    code_id = code_id or 1
 
     with contextlib.suppress((ValueError, TypeError)):
         code_id = int(code_id)
 
     if code_id not in _CODES:
-        raise GeneticCodeError(f"Unknown genetic code {code_id}")
+        raise GeneticCodeError(f"Unknown genetic code {code_id!r}")
     return _CODES[code_id]
 
 
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index 80bd29cb8..2b3a66dc8 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -1934,12 +1934,12 @@ class NucleicAcidSequenceMixin:
 
         # since we are realising the view, reverse complementing will be
         # dealt with, so rc=False
-        pep = gc.translate(array(seq), rc=False)
+        pep = gc.translate(array(seq), rc=False, incomplete_ok=incomplete_ok)
 
         if not include_stop and "*" in pep:
             raise new_alphabet.AlphabetError("stop codon in translation")
 
-        if not incomplete_ok and ("-" in pep or "X" in pep):
+        if not incomplete_ok and "X" in pep:
             raise new_alphabet.AlphabetError(
                 "Incomplete codon in translation, set incomplete_ok=True to "
                 "allow translation"


Test patch:
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index efe81e9c0..c821f839f 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -550,7 +550,7 @@ def test_deserialise_alphas(alpha):
 @pytest.fixture
 def calpha():
     gc = new_genetic_code.get_code(1)
-    return new_alphabet.CodonAlphabet(
+    return new_alphabet.SenseCodonAlphabet(
         words=gc.sense_codons, monomers=gc.moltype.alphabet
     )
 
@@ -570,7 +570,7 @@ def test_to_indices(calpha, seq):
 
 
 def test_codon_alphabet_invalid_codon(calpha):
-    with pytest.raises(ValueError):
+    with pytest.raises(new_alphabet.AlphabetError):
         calpha.to_index("TGA")
 
 
@@ -592,15 +592,41 @@ def test_codon_alphabet_not_is_valid(seq, calpha):
     assert not calpha.is_valid(seq)
 
 
+@pytest.mark.parametrize("as_array", (False, True))
+def test_codon_alphabet_to_indices(calpha, as_array):
+    seq = calpha.monomers.to_indices("GGGAAG") if as_array else "GGGAAG"
+    got = calpha.to_indices(seq)
+    expect = numpy.array(
+        [calpha.to_index("GGG"), calpha.to_index("AAG")], dtype=numpy.uint8
+    )
+    assert_allclose(got, expect)
+    assert isinstance(got, numpy.ndarray)
+    assert got.dtype == numpy.uint8
+
+
 def test_codon_alphabet_with_gap_motif(calpha):
     with_gap = calpha.with_gap_motif()
     assert len(with_gap) == len(calpha) + 1
+    assert with_gap.gap_char == "---"
+    # on the standard genetic code, gap index is 61
+def test_codon_alphabet_with_gap_motif_translation(calpha):
+    with_gap = calpha.with_gap_motif()
+    seq = "ATG---TAC"
+    translated = with_gap.to_indices(seq)
+    expected = [with_gap.to_index("ATG"), with_gap.gap_index, with_gap.to_index("TAC")]
+    assert_allclose(translated, expected)
+
+
+def test_codon_alphabet_missing(calpha):
+    # not provided on codon alphabet
+    assert calpha.missing_char is None
+    assert calpha.missing_index is None
 
 
 def test_codon_alphabet_serlialise_round_trip(calpha):
     from cogent3.util.deserialise import deserialise_object
 
     got = deserialise_object(calpha.to_json())
-    assert isinstance(got, new_alphabet.CodonAlphabet)
+    assert isinstance(got, new_alphabet.SenseCodonAlphabet)
     assert list(got) == list(calpha)
     assert calpha.to_index("TTC") == 1
diff --git a/tests/test_core/test_new_genetic_code.py b/tests/test_core/test_new_genetic_code.py
index d76aa8d87..36498277f 100644
--- a/tests/test_core/test_new_genetic_code.py
+++ b/tests/test_core/test_new_genetic_code.py
@@ -1,6 +1,6 @@
 import pytest
 
-from cogent3.core import new_genetic_code
+from cogent3.core import new_alphabet, new_genetic_code
 
 
 @pytest.mark.parametrize("gc", (1, "Standard", new_genetic_code.DEFAULT))
@@ -44,6 +44,12 @@ def test_sense_codons():
     assert len(code.sense_codons) == 61
 
 
+@pytest.mark.parametrize("invalid", ("AT", "TAAA", 23))
+def test_invalid_index(invalid):
+    with pytest.raises(new_genetic_code.InvalidCodonError):
+        new_genetic_code.DEFAULT[invalid]
+
+
 def test_translate():
     code = new_genetic_code.DEFAULT
     assert code.translate("ATGGGG") == code["ATG"] + code["GGG"]
@@ -56,6 +62,17 @@ def test_translate_rc():
     assert code.translate(rc, rc=True) == code["ATG"] + code["GGG"]
 
 
+def test_translate_incomplete_ok():
+    code = new_genetic_code.DEFAULT
+    assert code.translate("ATGG--", incomplete_ok=True) == code["ATG"] + "-"
+
+
+def test_translate_incomplete_not_ok():
+    code = new_genetic_code.DEFAULT
+    with pytest.raises(new_alphabet.AlphabetError):
+        assert code.translate("ATGG--", incomplete_ok=False)
+
+
 def test_to_table():
     code = new_genetic_code.DEFAULT
     table = code.to_table()
@@ -146,3 +163,10 @@ def test_get_alphabet_with_stop():
 
     alpha_gap = gc.get_alphabet(include_gap=True, include_stop=True)
     assert len(alpha_gap) == 65
+
+
+@pytest.mark.parametrize("repr_method", ("__str__", "__repr__", "_repr_html_"))
+def test_code_repr(repr_method):
+    gc = new_genetic_code.get_code(1)
+    got = getattr(gc, repr_method)()
+    assert isinstance(got, str)