Repository: cogent3/cogent3
Pull request #2165 (2024-12-12 03:47:30)

MAINT: more edits for integration of new_type
## Summary by Sourcery

Integrate new_type by enhancing the MolType class with improved handling of ambiguities and missing data, and refactor the Alphabet class to support motif subsets. Update tests to cover new functionality and ensure robustness.

Enhancements:
- Refactor the MolType class to use frozenset for complements and ambiguities, improving immutability and consistency.

Tests:
- Add tests for the new get_subset method in the Alphabet class, ensuring correct functionality and error handling.
- Update test cases to reflect changes in handling of missing and gap characters in the MolType class.


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index f59198dee..3454cfddc 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -18,6 +18,7 @@ StrORBytesORArray = typing.Union[str, bytes, numpy.ndarray]
 OptInt = typing.Optional[int]
 OptStr = typing.Optional[str]
 OptBytes = typing.Optional[bytes]
+PySeqStrOrBytes = typing.Sequence[str | bytes]
 
 
 @functools.singledispatch
@@ -277,6 +278,10 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
             assert _coerce_to_type(chars[0], missing) in chars
 
         consistent_words(chars, length=1)
+        if isinstance(chars, bytes):
+            # we need to convert to tuple in a way the preserves elements
+            # as bytes
+            chars = tuple(bytes([c]) for c in chars)
         return tuple.__new__(cls, chars, gap=gap, missing=missing)
 
     def __init__(
@@ -450,8 +455,8 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
 
     def as_bytes(self) -> bytes:
         """returns self as a byte string"""
-        if not isinstance(self[0], str):
-            return bytes(bytearray(self))
+        if isinstance(self[0], bytes):
+            return b"".join(self)
 
         return "".join(self).encode("utf8")
 
@@ -491,6 +496,29 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
     def get_word_alphabet(self, k: int, include_gap: bool = True) -> "KmerAlphabet":
         return self.get_kmer_alphabet(k, include_gap=include_gap)
 
+    def get_subset(
+        self,
+        motif_subset: PySeqStrOrBytes,
+        excluded: bool = False,
+    ) -> typing_extensions.Self:
+        """Returns a new Alphabet object containing a subset of motifs in self.
+
+        Raises an exception if any of the items in the subset are not already
+        in self.
+        """
+        self_set = set(self)
+        self_set |= {self.gap_char, self.missing_char}
+        self_set.discard(None)
+
+        if diff := set(motif_subset) - self_set:
+            raise AlphabetError(f"{diff!r} not members in self")
+
+        if excluded:
+            motif_subset = [m for m in self if m not in motif_subset]
+        gap = self.gap_char if self.gap_char in motif_subset else None
+        missing = self.missing_char if self.missing_char in motif_subset else None
+        return self.__class__(tuple(motif_subset), gap=gap, missing=missing)
+
 
 @register_deserialiser(get_object_provenance(CharAlphabet))
 def deserialise_char_alphabet(data: dict) -> CharAlphabet:
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index 41cf81fa9..2c3eec936 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -148,6 +148,22 @@ RNA_EXTENDED_PAIRS = {
 }
 
 
+def ambigs_and_missing(
+    ambiguities: dict[str, frozenset[str]],
+    gap_char: str | None,
+) -> dict[str, frozenset[str]]:
+    """returns new ambiguity dict with IUPAC missing as a key to most general code
+
+    Notes
+    -----
+    Includes the gap character if gap is not None
+    """
+    val = max(ambiguities.values(), key=len)
+    if gap_char:
+        val = val | frozenset((gap_char,))
+    return {**ambiguities, IUPAC_missing: val}
+
+
 class MolTypeError(TypeError): ...
 
 
@@ -403,8 +419,8 @@ class MolType:
     make_seq: dataclasses.InitVar[type]
     gap: OptStr = IUPAC_gap
     missing: OptStr = IUPAC_missing
-    complements: dataclasses.InitVar[dict[str, str] | None] = None
-    ambiguities: dict[str, tuple[str, ...]] | None = None
+    complements: dataclasses.InitVar[dict[str, frozenset[str]] | None] = None
+    ambiguities: dict[str, frozenset[str]] | None = None
     colors: dataclasses.InitVar[dict[str, str] | None] = None
     pairing_rules: dict[str, dict[frozenset[str], bool]] | None = None
     mw_calculator: WeightCalculator | None = None
@@ -427,7 +443,7 @@ class MolType:
         self,
         monomers: StrORBytes,
         make_seq: type,
-        complements: dict[str, str] | None,
+        complements: dict[str, frozenset[str]] | None,
         colors: dict[str, str] | None,
     ):
         self._colors = colors or defaultdict(_DefaultValue("black"))
@@ -483,6 +499,11 @@ class MolType:
             else None
         )
 
+        if self.missing and self.ambiguities:
+            # this MUST occur after the creation of the different alphabets otherwise
+            # missing data is added twice into an alphabet
+            self.ambiguities = ambigs_and_missing(self.ambiguities, self.gap)
+
         if complements:
             # assume we have a nucleic acid moltype
             dest = "".join(complements[c] for c in self.degen_gapped_alphabet)
@@ -1106,6 +1127,9 @@ class MolType:
 
     def degenerate_from_seq(self, seq: str) -> str:
         """Returns least degenerate symbol that encompasses a set of characters"""
+        if not self.ambiguities:
+            return self.missing
+
         symbols = frozenset(seq)
 
         # if the length of the set is 1, return the single character
@@ -1343,7 +1367,6 @@ ASCII = MolType(
     monomers="".join(ascii_letters),
     name="text",
     make_seq=new_sequence.Sequence,
-    missing=IUPAC_missing,
 )
 
 DNA = MolType(
@@ -1391,9 +1414,9 @@ BYTES = MolType(
     # want to prematurely assume _anything_ about the data.
     monomers=bytes(bytearray(range(2**8))),
     name="bytes",
+    make_seq=new_sequence.ByteSequence,
     gap=None,
     missing=None,
-    make_seq=new_sequence.ByteSequence,
 )
 
 # the None value catches cases where a moltype has no label attribute


Test patch:
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index 61f4d6350..c049f86ef 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -665,12 +665,12 @@ def test_codon_alphabet_serlialise_round_trip(calpha):
 
 
 @pytest.fixture(params=new_moltype.DNA.iter_alphabets())
-def alpha(request):
+def dna_alpha(request):
     return request.param
 
 
-def test_alphabet_moltype(alpha):
-    assert alpha.moltype is new_moltype.DNA
+def test_alphabet_moltype(dna_alpha):
+    assert dna_alpha.moltype is new_moltype.DNA
 
 
 def test_alpha_no_moltype():
@@ -720,3 +720,23 @@ def test_sensecodon_alphabet_to_indices_list_tuple(cast, calpha):
         dtype=calpha.dtype,
     )
     assert_allclose(got, expect)
+
+
+def test_get_subset():
+    dna_alpha = new_moltype.get_moltype("dna").gapped_alphabet
+    got = dna_alpha.get_subset(("A", "G"))
+    assert len(got) == 2
+    assert got.gap_char is None
+    assert got.missing_char is None
+    got = dna_alpha.get_subset(("A", "G", "-"))
+    assert got.gap_char == "-"
+    assert len(got) == 3
+    got = dna_alpha.get_subset("-", excluded=True)
+    assert got.gap_char is None
+    assert len(got) == 4
+
+
+def test_get_subset_invalid():
+    dna_alpha = new_moltype.get_moltype("dna").alphabet
+    with pytest.raises(new_alphabet.AlphabetError):
+        dna_alpha.get_subset(("A", "G", "-"))
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index 3d36ce38a..fc8f569fa 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -373,8 +373,10 @@ def test_most_degenerate_alphabet(label):
     got = moltype.most_degen_alphabet()
     # expected value for number of characters is
     # length of monomers + len(moltype gaps) + len(ambiguities)
-    num_ambigs = len(moltype.ambiguities or [])
-    expected = len(moltype.alphabet) + len(moltype.gaps or []) + num_ambigs
+    ambigs = set(moltype.ambiguities or ())
+    ambigs |= set(moltype.missing) if moltype.missing else set()
+    num_ambigs = len(ambigs)
+    expected = len(moltype.alphabet) + len(moltype.gap or "") + num_ambigs
     assert len(got) == expected
 
 
@@ -572,6 +574,11 @@ def test_degenerate_from_seq():
     assert p("AC-") == "?"
 
 
+def test_degenerate_from_seq_general():
+    dna = new_moltype.get_moltype("dna")
+    assert dna.degenerate_from_seq("ACGT--") == "?"
+
+
 @pytest.mark.parametrize(
     "moltype,zero,second,last",
     (