Repository: cogent3/cogent3
Pull request #2085 (2024-10-31 23:00:26)

ENH: add to_regex to new MolType and GeneticCode
## Summary by Sourcery

Add `to_regex` methods to the `GeneticCode` and `MolType` classes to generate regex patterns from sequences, and implement corresponding tests to ensure functionality.

New Features:
- Introduce a `to_regex` method in the `GeneticCode` class to convert amino acid sequences into regex patterns that match corresponding DNA sequences.
- Add a `to_regex` method to the `MolType` class to expand sequences with ambiguities into regex patterns.

Tests:
- Add tests for the `to_regex` method in the `GeneticCode` class to ensure it correctly creates regex patterns from amino acid sequences.
- Add tests for the `to_regex` method in the `MolType` class to validate regex pattern generation and handle invalid sequences.


Patch:
diff --git a/src/cogent3/core/new_genetic_code.py b/src/cogent3/core/new_genetic_code.py
index b3253d307..67992a18b 100644
--- a/src/cogent3/core/new_genetic_code.py
+++ b/src/cogent3/core/new_genetic_code.py
@@ -280,6 +280,29 @@ class GeneticCode:
             words=codons, monomers=self.moltype.alphabet, gap=gap
         )
 
+    def to_regex(self, seq: typing.Union[str, "Sequence"]) -> str:
+        """returns a regex pattern with an amino acid expanded to its codon set
+
+        Parameters
+        ----------
+        seq
+            a Sequence or string of amino acids
+        """
+        from .new_moltype import PROTEIN_WITH_STOP_ambiguities as ambigs
+
+        seq = list(str(seq))
+        mappings = []
+        for aa in seq:
+            aa = ambigs[aa] if aa in ambigs else [aa]
+            codons = []
+            for a in aa:
+                codons.extend(self[a])
+
+            # we create a regex non-capturing group for each amino acid
+            mappings.append(f"(?:{'|'.join(codons)})")
+
+        return "".join(mappings)
+
 
 _mapping_cols = "ncbi_code_sequence", "ID", "name", "ncbi_start_codon_map"
 # code mappings are based on the product of bases in order TCAG
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index dc78a0a7a..2c2005698 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -1174,6 +1174,18 @@ class MolType:
         """returns the most degenerate alphabet for this instance"""
         return next(self.iter_alphabets())
 
+    def to_regex(self, seq: str) -> str:
+        """returns a regex pattern with ambiguities expanded to a character set"""
+        if not self.is_valid(seq):
+            raise ValueError(f"'{seq}' is invalid for this moltype")
+
+        degen_indices = self.get_degenerate_positions(seq=seq, include_gap=False)
+        seq = list(seq)  # seq can now be modified
+        for index in degen_indices:
+            expanded = self.ambiguities[seq[index]]
+            seq[index] = f"[{''.join(sorted(expanded))}]"
+        return "".join(seq)
+
 
 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_genetic_code.py b/tests/test_core/test_new_genetic_code.py
index d76aa8d87..9bab8be5b 100644
--- a/tests/test_core/test_new_genetic_code.py
+++ b/tests/test_core/test_new_genetic_code.py
@@ -146,3 +146,22 @@ def test_get_alphabet_with_stop():
 
     alpha_gap = gc.get_alphabet(include_gap=True, include_stop=True)
     assert len(alpha_gap) == 65
+
+
+def test_to_regex():
+    """creates a regex from aa seq to match a DNA sequence"""
+    import re
+
+    from cogent3 import make_seq
+
+    dna = "ACCGAACAGGGC"
+    aa = "TEQG"
+    pattern = new_genetic_code.DEFAULT.to_regex(aa)
+    assert "".join(re.findall(pattern, dna)) == dna
+    # note that Z is Q or E
+    aa = "TZQG"
+    pattern = new_genetic_code.DEFAULT.to_regex(aa)
+    assert "".join(re.findall(pattern, dna)) == dna
+    aa = make_seq(aa, moltype="protein", new_type=True)
+    pattern = new_genetic_code.DEFAULT.to_regex(aa)
+    assert "".join(re.findall(pattern, dna)) == dna
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index a2bbd6217..b311525e8 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -596,3 +596,13 @@ def test_make_seq_with_seq_invalid_moltype():
     seq = new_moltype.PROTEIN.make_seq(seq="ACQK")
     with pytest.raises(ValueError):
         _ = new_moltype.DNA.make_seq(seq=seq)
+
+
+def test_to_regex():
+    """returns a valid regex"""
+    seq = "ACYGR"
+    regular_expression = new_moltype.DNA.to_regex(seq=seq)
+    assert regular_expression == "AC[CT]G[AG]"
+    # raises an exception if a string is already a regex, or invalid
+    with pytest.raises(ValueError):
+        new_moltype.DNA.to_regex("(?:GAT|GAC)(?:GGT|GGC|GGA|GGG)(?:GAT|GAC)(?:CAA|CAG)")