Repository: cogent3/cogent3
Pull request #1437 (2023-05-13 00:23:32)

Mods to sequences and alignments


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index d2fb28693..dbef886c1 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -1326,7 +1326,9 @@ class _SequenceCollectionBase:
         """len of SequenceCollection returns length of longest sequence."""
         return self.seq_len
 
-    def get_translation(self, gc=None, incomplete_ok=False, **kwargs):
+    def get_translation(
+        self, gc=None, incomplete_ok=False, include_stop=False, **kwargs
+    ):
         """translate from nucleic acid to protein
 
         Parameters
@@ -1334,9 +1336,11 @@ class _SequenceCollectionBase:
         gc
             genetic code, either the number or name
             (use cogent3.core.genetic_code.available_codes)
-        incomplete_ok : bool
+        incomplete_ok
             codons that are mixes of nucleotide and gaps converted to '?'.
             raises a ValueError if False
+        include_stop
+            whether to allow a stops in the translated sequence
         kwargs
             related to construction of the resulting object
 
@@ -1344,21 +1348,22 @@ class _SequenceCollectionBase:
         -------
         A new instance of self translated into protein
         """
+        if len(self.moltype.alphabet) != 4:
+            raise TypeError("Must be a DNA/RNA")
+
         translated = []
-        aligned = isinstance(self, Alignment)
-        kwargs["moltype"] = cogent3.PROTEIN
         # do the translation
-        try:
-            for seqname in self.names:
-                if aligned:
-                    seq = self.get_gapped_seq(seqname)
-                else:
-                    seq = self.named_seqs[seqname]
-                pep = seq.get_translation(gc, incomplete_ok=incomplete_ok)
-                translated.append((seqname, pep))
-            return self.__class__(translated, info=self.info, **kwargs)
-        except AttributeError as msg:
-            raise AttributeError(f"{msg} -- {'Did you set a DNA moltype?'}")
+        for seqname in self.names:
+            try:
+                seq = self.get_gapped_seq(seqname)
+            except AttributeError:
+                seq = self.named_seqs[seqname]
+            pep = seq.get_translation(
+                gc, incomplete_ok=incomplete_ok, include_stop=include_stop
+            )
+            translated.append((seqname, pep))
+        kwargs["moltype"] = pep.moltype
+        return self.__class__(translated, info=self.info, **kwargs)
 
     def get_seq(self, seqname):
         """Return a sequence object for the specified seqname."""
diff --git a/src/cogent3/core/genetic_code.py b/src/cogent3/core/genetic_code.py
index 111f6d369..5ad41f6c5 100644
--- a/src/cogent3/core/genetic_code.py
+++ b/src/cogent3/core/genetic_code.py
@@ -286,7 +286,7 @@ class GeneticCode:
 
     def is_start(self, codon):
         """Returns True if codon is a start codon, False otherwise."""
-        fixed_codon = codon.upper().replace("U", "T")
+        fixed_codon = str(codon).upper().replace("U", "T")
         return fixed_codon in self.start_codons
 
     def is_stop(self, codon):
diff --git a/src/cogent3/core/sequence.py b/src/cogent3/core/sequence.py
index fde467fc7..aa19a252f 100644
--- a/src/cogent3/core/sequence.py
+++ b/src/cogent3/core/sequence.py
@@ -10,6 +10,8 @@ Sequences are intended to be immutable. This is not enforced by the code for
 performance reasons, but don't alter the MolType or the sequence data after
 creation.
 """
+from __future__ import annotations
+
 import contextlib
 import copy
 import json
@@ -1463,6 +1465,7 @@ class Sequence(SequenceI):
         ]
 
     def parse_out_gaps(self):
+        """returns Map corresponding to gap locations and ungapped Sequence"""
         gapless = []
         segments = []
         nongap = re.compile(f"([^{re.escape('-')}]+)")
@@ -1474,7 +1477,7 @@ class Sequence(SequenceI):
             "".join(gapless), name=self.get_name(), info=self.info, preserve_case=True
         )
         seq.annotation_db = copy.deepcopy(self.annotation_db)
-        return (map, seq)
+        return map, seq
 
     def replace(self, oldchar, newchar):
         """return new instance with oldchar replaced by newchar"""
@@ -1669,14 +1672,23 @@ class NucleicAcidSequence(Sequence):
         ----------
         gc
             name or ID of genetic code
-        incomplete_ok : bool
+        incomplete_ok
             codons that are mixes of nucleotide and gaps converted to '?'.
             raises a ValueError if False
+        include_stop
+            allows stop codons in translation
 
         Returns
         -------
         sequence of PROTEIN moltype
+
+        Raises
+        ------
+        AlphabetError if include_stop is False and a stop codon occurs
         """
+        from cogent3.core.moltype import get_moltype
+
+        protein = get_moltype("protein_with_stop" if include_stop else "protein")
         gc = get_code(gc)
         codon_alphabet = gc.get_alphabet(include_stop=include_stop).with_gap_motif()
         # translate the codons
@@ -1699,36 +1711,15 @@ class NucleicAcidSequence(Sequence):
                         raise AlphabetError(f"incomplete codon {codon} in {self.name}")
                 else:
                     aa = gc[codon]
-                    if aa == "*":
+                    if aa == "*" and not include_stop:
                         continue
                 trans.append(aa)
             if not trans:
                 raise ValueError(orig_codon)
-            aa = self.protein.what_ambiguity(trans)
+            aa = protein.what_ambiguity(trans)
             translation.append(aa)
 
-        translation = self.protein.make_seq(seq="".join(translation), name=self.name)
-
-        return translation
-
-    def get_orf_positions(self, gc=None, atg=False):
-        gc = get_code(gc)
-        orfs = []
-        start = None
-        protein = self.get_translation(gc=gc)
-        for posn, aa in enumerate(protein):
-            posn *= 3
-            if aa == "*":
-                if start is not None:
-                    orfs.append((start, posn))
-                start = None
-            else:
-                if start is None:
-                    if (not atg) or gc.is_start(self[posn : posn + 3]):
-                        start = posn
-        if start is not None:
-            orfs.append((start, posn + 3))
-        return orfs
+        return protein.make_seq(seq="".join(translation), name=self.name)
 
     def to_rna(self):
         """Returns copy of self as RNA."""
diff --git a/src/cogent3/parse/genbank.py b/src/cogent3/parse/genbank.py
index 555636686..10a651266 100644
--- a/src/cogent3/parse/genbank.py
+++ b/src/cogent3/parse/genbank.py
@@ -696,10 +696,10 @@ def RichGenbankParser(
         a GenbankAnnotationDb instance to which feature data will be
         added
     """
-    info_excludes = info_excludes or []
+    info_excludes = info_excludes or ["sequence", "features"]
     moltype = get_moltype(moltype) if moltype else None
     for rec in MinimalGenbankParser(handle):
-        info = Info()
+        info = {}
         # populate the info object, excluding the sequence
         for label, value in list(rec.items()):
             if label in info_excludes:
@@ -715,6 +715,7 @@ def RichGenbankParser(
         else:
             rec_moltype = moltype
 
+        info = Info(genbank_record=info)
         try:
             seq = rec_moltype.make_seq(
                 rec["sequence"].upper(), info=info, name=rec["locus"]
@@ -731,6 +732,6 @@ def RichGenbankParser(
 
         db = getattr(db, "db", None)
         seq.annotation_db = GenbankAnnotationDb(
-            data=rec["features"], seqid=rec["locus"], db=db
+            data=rec.pop("features", None), seqid=rec["locus"], db=db
         )
-        yield (rec["locus"], seq)
+        yield rec["locus"], seq


Test patch:
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index 28c9c0ec3..92b62601c 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -894,23 +894,6 @@ class SequenceCollectionBaseTests(object):
         aln = self.Class([("a", "AAAA"), ("b", "TTTT"), ("c", "CCCC")])
         self.assertEqual(len(aln), 4)
 
-    def test_get_translation(self):
-        """SequenceCollection.get_translation translates each seq"""
-        for seqs in [
-            {"seq1": "GATTTT", "seq2": "GATC??"},
-            {"seq1": "GAT---", "seq2": "?GATCT"},
-        ]:
-            alignment = self.Class(data=seqs, moltype=DNA)
-            got = alignment.get_translation()
-            self.assertEqual(len(got), 2)
-            self.assertEqual(got.moltype, PROTEIN)
-            # check for a failure when no moltype specified
-            alignment = self.Class(data=seqs)
-            try:
-                alignment.get_translation()
-            except AttributeError:
-                pass
-
     def test_get_translation_info(self):
         """SequenceCollection.get_translation preserves info attribute"""
         for seqs in [
@@ -3455,3 +3438,36 @@ def test_distance_matrix_passes_correct_moltype(moltype):
     data = [("s1", "ACGTA"), ("s2", "ACGTA")]
     seqs = make_unaligned_seqs(data=data, moltype=moltype)
     seqs.distance_matrix()
+
+
+@pytest.mark.parametrize("cls", (SequenceCollection, Alignment, ArrayAlignment))
+@pytest.mark.parametrize(
+    "seqs", ({"seq1": "GATTTT", "seq2": "GATC??"}, {"seq1": "GAT---", "seq2": "?GATCT"})
+)
+def test_get_translation2(cls, seqs):
+    """SequenceCollection.get_translation translates each seq"""
+    alignment = cls(data=seqs, moltype=DNA)
+    got = alignment.get_translation()
+    assert len(got) == 2
+    assert got.moltype == PROTEIN
+
+
+@pytest.mark.parametrize("cls", (SequenceCollection, Alignment, ArrayAlignment))
+def test_get_translation_with_stop(cls):
+    """SequenceCollection.get_translation translates each seq"""
+    seqs = {"seq1": "GATTAG", "seq2": "?GATCT"}
+    alignment = cls(data=seqs, moltype=DNA)
+    got = alignment.get_translation(include_stop=True)
+    assert got.to_dict() == {"seq1": "D*", "seq2": "XS"}
+
+
+@pytest.mark.parametrize("cls", (SequenceCollection, Alignment, ArrayAlignment))
+@pytest.mark.parametrize(
+    "seqs", ({"seq1": "GATTTT", "seq2": "GATC??"}, {"seq1": "GAT---", "seq2": "?GATCT"})
+)
+def test_get_translation_error(cls, seqs):
+    """SequenceCollection.get_translation translates each seq"""
+    # check for a failure when no moltype specified
+    alignment = cls(data=seqs)
+    with pytest.raises(TypeError):
+        alignment.get_translation()
diff --git a/tests/test_core/test_core_standalone.py b/tests/test_core/test_core_standalone.py
index 85b04ef38..c9fcf052a 100644
--- a/tests/test_core/test_core_standalone.py
+++ b/tests/test_core/test_core_standalone.py
@@ -669,20 +669,6 @@ class AlignmentTestMethods(unittest.TestCase):
         for char in seqs[0]:
             self.assertEqual(seqs[0].count(char), 2)
 
-    def test_translate(self):
-        for seqs in [
-            {"seq1": "GATTTT", "seq2": "GATC??"},
-            {"seq1": "GAT---", "seq2": "?GATCT"},
-        ]:
-            alignment = make_aligned_seqs(data=seqs, moltype=DNA)
-            self.assertEqual(len(alignment.get_translation()), 2)
-            # check for a failure when no moltype specified
-            alignment = make_aligned_seqs(data=seqs)
-            try:
-                alignment.get_translation()
-            except AttributeError:
-                pass
-
     def test_seqnames(self):
         s1 = self.alignment.get_seq("Mouse")
         self.assertEqual(s1.get_name(), "Mouse")
diff --git a/tests/test_core/test_sequence.py b/tests/test_core/test_sequence.py
index fafc8d294..4e44c1d36 100644
--- a/tests/test_core/test_sequence.py
+++ b/tests/test_core/test_sequence.py
@@ -1259,6 +1259,12 @@ def test_translation(seq, rc):
     assert str(aa) == "MF"
 
 
+def test_get_translation_include_stop():
+    s = DNA.make_seq("ATTTAACTT", name="s1")
+    aa = s.get_translation(include_stop=True)
+    assert str(aa) == "I*L"
+
+
 @pytest.mark.parametrize("start", (None, 0, 1, 10, -1, -10))
 @pytest.mark.parametrize("stop", (None, 10, 8, 1, 0, -1, -11))
 @pytest.mark.parametrize("step", (None, 1, 2, -1, -2))
diff --git a/tests/test_parse/test_genbank.py b/tests/test_parse/test_genbank.py
index de9f8d134..50ce1a3bc 100644
--- a/tests/test_parse/test_genbank.py
+++ b/tests/test_parse/test_genbank.py
@@ -2,6 +2,8 @@
 """
 from unittest import TestCase
 
+import pytest
+
 from cogent3.parse.genbank import (
     Location,
     LocationList,
@@ -420,62 +422,6 @@ ORIGIN
             ],
         )
 
-    def test_rich_parser(self):
-        """correctly constructs +/- strand features"""
-        with open("data/annotated_seq.gb") as infile:
-            parser = RichGenbankParser(infile)
-            seq = [s for l, s in parser][0]
-
-        cds = dict([(f.name, f) for f in seq.get_features(biotype="CDS")])
-        expects = {
-            "CNA00110": "MAGYDARYGNPLDPMSGGRPSPPETSQQDAYEYSKHGSSSGYLGQLPLGAD"
-            "SAQAETASALRTLFGEGADVQALQEPPNQINTLAEGAAVAETGGVLGGDTTRSDNEALAIDPSL"
-            "SEQAAPAPKDSTETPDDRSRSPSSGNHHHHHPAVKRKATSRAGMLARGGACEFCKRRKLKCSAEL"
-            "PACANCVKSGKECVYAQKKQRSRVKVLEDRLQELEKRLEQGQAGAASASGGDSGAHAASSVYTAP"
-            "SLGSGGGSELTVEQTLVHNVDPSLLPPSEYDEAFILHDFDSFADMRKQETQLEPDLMTLADAAAA"
-            "DTPAAAAAETNDPWAKMSPEEIVKEIIKVATGGKGEGERIISHLVQTYMNSTVNTWHPLVIPPMD"
-            "LVSRVSRTTPDPIHPTLLLSLIPALLPLSPIQSLRHPAIPLLLLPHARAHSVQAITQSDPRVLDT"
-            "IIAGVSRAYSFFNEAKNIDGWVDCVAATSLVRAAGLTKQGGVGERFVPEDRVPAERLAKRRREAG"
-            "LRALMHKGAIVPPPESWYQFGQRVNLFWTSYICDRAAAIGWGWPSSYNDEDITTPWPKDDYKSVQ"
-            "ALLDDTTIHTFLSPLAPAPAPATPDSDLCAQAKSITLLYHAQRLLDSPPELSTPEKTHRLLGLTE"
-            "GYMESLEKMRGPRMRAGKLSSVWMILYTTIAVLHSKDGFDKCDPDGADQVSITRVVAAADKVLEL"
-            "VSAVQNTGDTHLSSCDVISSVLFLHLARLMIQYTNRLRLRVQDSALVSTLRAKTESFKRALIDQG"
-            "ERLVFAQVAAQMLENYHVGAEWKAGEWERADGGDWRGV",
-            "CNA00120": "MDFSQFNGAEQAHMSKVIEKKQMQDFMRLYSGLVEKCFNACAQD"
-            "FTSKALTTNETTCVQNCTDKFLKHSERVGARFAEHNAGMLSPYGAASLMASQSKCRAP"
-            "DSNGLGVFCKWRRIKSTVVLYNHLACIKQMDNRF",
-        }
-
-        for locus in cds:
-            got = cds[locus].get_slice().trim_stop_codon().get_translation()
-            self.assertEqual(str(got), expects[locus])
-
-    def test_rich_parser_moltype(self):
-        """correctly handles moltypes"""
-        with open("data/annotated_seq.gb") as infile:
-            parser = RichGenbankParser(infile)
-            got_1 = [s for _, s in parser][0]
-
-        # name formed from /product value
-        feature_ids = {"CNA00110", "CNA00120"}
-        got = {f.name for f in got_1.get_features(biotype="mRNA")}
-        self.assertEqual(got, {"CNA00110", "CNA00120"})
-
-        # the file defines itself as DNA
-        self.assertEqual(got_1.moltype.label, "dna")
-
-        # but that is overridden by user setting moltype explicitly
-        for moltype in ("dna", "rna", "text"):
-            with open("data/annotated_seq.gb") as infile:
-                parser = RichGenbankParser(infile, moltype=moltype)
-                got_2 = [s for _, s in parser][0]
-            self.assertEqual(
-                got_1.annotation_db.num_matches(), got_2.annotation_db.num_matches()
-            )
-            self.assertEqual(got_2.moltype.label, moltype)
-            got = {f.name for f in got_1.get_features(biotype="mRNA")}
-            self.assertEqual(got, feature_ids)
-
 
 class LocationTests(TestCase):
     """Tests of the Location class."""
@@ -547,3 +493,69 @@ def test_location_list_get_coordinates():
     g = parse_location_line(l)
     spans = g.get_coordinates()
     assert spans == [(5669, 5918), (5964, 6126)]
+
+
+@pytest.fixture(scope="session")
+def rich_gb():
+    with open("data/annotated_seq.gb") as infile:
+        parser = RichGenbankParser(infile)
+        seq = [s for l, s in parser][0]
+    return seq
+
+
+def test_rich_parser(rich_gb):
+    """correctly constructs +/- strand features"""
+    cds = dict([(f.name, f) for f in rich_gb.get_features(biotype="CDS")])
+    expects = {
+        "CNA00110": "MAGYDARYGNPLDPMSGGRPSPPETSQQDAYEYSKHGSSSGYLGQLPLGAD"
+        "SAQAETASALRTLFGEGADVQALQEPPNQINTLAEGAAVAETGGVLGGDTTRSDNEALAIDPSL"
+        "SEQAAPAPKDSTETPDDRSRSPSSGNHHHHHPAVKRKATSRAGMLARGGACEFCKRRKLKCSAEL"
+        "PACANCVKSGKECVYAQKKQRSRVKVLEDRLQELEKRLEQGQAGAASASGGDSGAHAASSVYTAP"
+        "SLGSGGGSELTVEQTLVHNVDPSLLPPSEYDEAFILHDFDSFADMRKQETQLEPDLMTLADAAAA"
+        "DTPAAAAAETNDPWAKMSPEEIVKEIIKVATGGKGEGERIISHLVQTYMNSTVNTWHPLVIPPMD"
+        "LVSRVSRTTPDPIHPTLLLSLIPALLPLSPIQSLRHPAIPLLLLPHARAHSVQAITQSDPRVLDT"
+        "IIAGVSRAYSFFNEAKNIDGWVDCVAATSLVRAAGLTKQGGVGERFVPEDRVPAERLAKRRREAG"
+        "LRALMHKGAIVPPPESWYQFGQRVNLFWTSYICDRAAAIGWGWPSSYNDEDITTPWPKDDYKSVQ"
+        "ALLDDTTIHTFLSPLAPAPAPATPDSDLCAQAKSITLLYHAQRLLDSPPELSTPEKTHRLLGLTE"
+        "GYMESLEKMRGPRMRAGKLSSVWMILYTTIAVLHSKDGFDKCDPDGADQVSITRVVAAADKVLEL"
+        "VSAVQNTGDTHLSSCDVISSVLFLHLARLMIQYTNRLRLRVQDSALVSTLRAKTESFKRALIDQG"
+        "ERLVFAQVAAQMLENYHVGAEWKAGEWERADGGDWRGV",
+        "CNA00120": "MDFSQFNGAEQAHMSKVIEKKQMQDFMRLYSGLVEKCFNACAQD"
+        "FTSKALTTNETTCVQNCTDKFLKHSERVGARFAEHNAGMLSPYGAASLMASQSKCRAP"
+        "DSNGLGVFCKWRRIKSTVVLYNHLACIKQMDNRF",
+    }
+
+    for locus in cds:
+        got = cds[locus].get_slice().trim_stop_codon().get_translation()
+        assert str(got) == expects[locus]
+
+
+def test_rich_parser_moltype(rich_gb):
+    """correctly handles moltypes"""
+
+    # name formed from /product value
+    feature_ids = {"CNA00110", "CNA00120"}
+    got = {f.name for f in rich_gb.get_features(biotype="mRNA")}
+    assert got == {"CNA00110", "CNA00120"}
+    # the file defines itself as DNA
+    assert rich_gb.moltype.label == "dna"
+    got = {f.name for f in rich_gb.get_features(biotype="mRNA")}
+    assert got == feature_ids
+
+
+@pytest.mark.parametrize("moltype", ("dna", "rna", "text"))
+def test_moltype_overrides(moltype, rich_gb):
+    # moltype is overridden by user setting moltype explicitly
+    with open("data/annotated_seq.gb") as infile:
+        parser = RichGenbankParser(infile, moltype=moltype)
+        got_2 = [s for _, s in parser][0]
+
+    assert rich_gb.annotation_db.num_matches() == got_2.annotation_db.num_matches()
+
+    assert got_2.moltype.label == moltype
+
+
+def test_rich_parser_info(rich_gb):
+    """seq.info stores genbank_record"""
+    assert "genbank_record" in rich_gb.info
+    assert rich_gb.info.genbank_record["locus"] == rich_gb.name