Repository: cogent3/cogent3
Pull request #2076 (2024-10-28 22:49:57)

DEV: fixes to naming of annotated sequences
## Summary by Sourcery

Fix naming inconsistencies in annotated sequences and enhance sequence loading to support Genbank format with improved annotation handling. Add tests for Genbank format and annotation verification.

Bug Fixes:
- Fix the naming of annotated sequences to ensure consistency across the codebase.

Enhancements:
- Enhance the sequence loading functionality to support Genbank format and handle annotations more effectively.

Tests:
- Add tests for loading sequences with Genbank format and verifying annotation handling.
- Introduce parameterized tests for checking Genbank format recognition.


Patch:
diff --git a/src/cogent3/__init__.py b/src/cogent3/__init__.py
index 1b8e1b40a..4acd6580c 100644
--- a/src/cogent3/__init__.py
+++ b/src/cogent3/__init__.py
@@ -11,13 +11,13 @@ from typing import Callable, Optional, Union
 
 from cogent3._version import __version__
 from cogent3.app import app_help, available_apps, get_app, open_data_store
+from cogent3.core import annotation_db as _anno_db
 from cogent3.core.alignment import (
     Alignment,
     ArrayAlignment,
     Sequence,
     SequenceCollection,
 )
-from cogent3.core.annotation_db import load_annotations
 from cogent3.core.genetic_code import available_codes, get_code
 
 # note that moltype has to be imported last, because it sets the moltype in
@@ -35,7 +35,7 @@ from cogent3.evolve.fast_distance import available_distances, get_distance_calcu
 from cogent3.evolve.models import available_models, get_model
 from cogent3.parse.cogent3_json import load_from_json
 from cogent3.parse.newick import parse_string as newick_parse_string
-from cogent3.parse.sequence import get_parser
+from cogent3.parse.sequence import get_parser, is_genbank
 from cogent3.parse.table import load_delimited
 from cogent3.parse.tree_xml import parse_string as tree_xml_parse_string
 from cogent3.util.io import get_format_suffixes, open_
@@ -72,6 +72,8 @@ import logging
 __numba_logger = logging.getLogger("numba")
 __numba_logger.setLevel(logging.WARNING)
 
+load_annotations = _anno_db.load_annotations
+
 
 def make_seq(
     seq,
@@ -79,6 +81,7 @@ def make_seq(
     moltype=None,
     new_type: bool = False,
     annotation_offset: int = 0,
+    annotation_db: Optional[_anno_db.SupportsFeatures] = None,
     **kw: dict,
 ):  # refactor: type hinting, need to capture optional args and the return type
     """
@@ -113,6 +116,8 @@ def make_seq(
     seq = moltype.make_seq(
         seq=seq, name=name, annotation_offset=annotation_offset, **kw
     )
+    if annotation_db:
+        seq.annotation_db = annotation_db
     return seq
 
 
@@ -275,6 +280,30 @@ def _load_seqs(file_format, filename, fmt, parser_kw):
     return list(parser(filename, **parser_kw))
 
 
+T = Optional[_anno_db.SupportsFeatures]
+
+
+def _load_genbank_seq(
+    filename: os.PathLike, parser_kw: dict, just_seq: bool = False
+) -> tuple[str, str, T]:
+    """utility function for loading sequences"""
+    from cogent3.parse.genbank import iter_genbank_records
+
+    for name, seq, features in iter_genbank_records(filename, **parser_kw):
+        break
+    else:
+        raise ValueError(f"No sequences found in {filename}")
+
+    db = (
+        None
+        if just_seq
+        else _anno_db.GenbankAnnotationDb(
+            data=features.pop("features", None), seqid=name
+        )
+    )
+    return name, seq, db
+
+
 def load_seq(
     filename: os.PathLike,
     annotation_path: Optional[os.PathLike] = None,
@@ -292,17 +321,19 @@ def load_seq(
 
     Parameters
     ----------
-    filename : str
+    filename
         path to sequence file
-    format : str
+    annotation_path
+        path to annotation file, ignored if format is genbank
+    format
         sequence file format, if not specified tries to guess from the path suffix
-    moltype : str
+    moltype
         the moltype, eg DNA, PROTEIN, 'dna', 'protein'
-    label_to_name : callable
+    label_to_name
         function for converting original name into another name.
-    parser_kw : dict
+    parser_kw
         optional arguments for the parser
-    info : dict
+    info
         a dict from which to make an info object
     new_type
         if True, returns a new type Sequence (cogent3.core.new_sequence.Sequence)
@@ -326,12 +357,22 @@ def load_seq(
     info["source"] = str(filename)
     file_format, _ = get_format_suffixes(filename)
     if file_format == "json":
-        seq = load_from_json(filename, (Sequence,))
+        seq = load_from_json(filename, (Sequence,))  # need to support new seq here
         seq.name = label_to_name(seq.name) if label_to_name else seq.name
         return seq
 
-    data = _load_seqs(file_format, filename, format, parser_kw)
-    name, seq = data[0]
+    if is_genbank(format or file_format):
+        name, seq, db = _load_genbank_seq(
+            filename, parser_kw or {}, just_seq=annotation_path is not None
+        )
+    else:
+        db = None
+        data = _load_seqs(file_format, filename, format, parser_kw)
+        name, seq = data[0]
+
+    if annotation_path is not None:
+        db = load_annotations(path=annotation_path, seqids=[name])
+
     name = label_to_name(name) if label_to_name else name
 
     result = make_seq(
@@ -340,15 +381,11 @@ def load_seq(
         moltype=moltype,
         new_type=new_type,
         annotation_offset=annotation_offset,
+        annotation_db=db,
         **kw,
     )
     result.info.update(info)
 
-    if getattr(seq, "annotation_db", None):
-        result.annotation_db = seq.annotation_db
-
-    if annotation_path is not None:
-        result.annotation_db = load_annotations(path=annotation_path, seqids=[name])
     return result
 
 
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index df635b611..c33f9be83 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -530,7 +530,9 @@ class SequenceCollection:
         # seqview is given the name of the parent (if different from the current name)
         # the sequence is given the current name
         sv = self._seqs_data.get_view(self._name_map.get(name, name))
-        return self.moltype.make_seq(seq=sv, name=name)
+        seq = self.moltype.make_seq(seq=sv, name=name)
+        seq.replace_annotation_db(self.annotation_db)
+        return seq
 
     @property
     def seqs(self) -> _IndexableSeqs:
@@ -807,7 +809,7 @@ class SequenceCollection:
         """
         seqs = {
             name: self.moltype.degap(self._seqs_data.get_seq_array(seqid=name))
-            for name in self.names
+            for name in self._name_map.values()
         }
         seqs_data = self._seqs_data.from_seqs(
             data=coerce_to_seqs_data_dict(seqs),
@@ -916,7 +918,7 @@ class SequenceCollection:
         seqs_data = coerce_to_seqs_data_dict(translated, label_to_name=None)
         pep_moltype = pep.moltype
 
-        seqs_data = self._seqs_data.__class__(
+        seqs_data = self._seqs_data.from_seqs(
             data=seqs_data,
             alphabet=pep_moltype.most_degen_alphabet(),
             offset=self._seqs_data.offset,
@@ -1699,12 +1701,12 @@ class SequenceCollection:
             )
 
         new_seqs = {}
-        for seq_name in self.names:
+        for seq_name in self._name_map.values():
             seq = self._seqs_data.get_seq_str(seqid=seq_name)
             padded_seq = seq + "-" * (pad_length - len(seq))
             new_seqs[seq_name] = padded_seq
 
-        seqs_data = self._seqs_data.__class__(
+        seqs_data = self._seqs_data.from_seqs(
             data=new_seqs,
             alphabet=self._seqs_data.alphabet,
             offset=self._seqs_data.offset,
@@ -2272,6 +2274,7 @@ def make_unaligned_seqs(
     source: OptPathType = None,
     annotation_db: SupportsFeatures = None,
     offset: dict[str, int] = None,
+    name_map: dict[str, str] = None,
 ) -> SequenceCollection:
     """Initialise an unaligned collection of sequences.
 
@@ -2315,6 +2318,18 @@ def make_unaligned_seqs(
 
     annotation_db = annotation_db or merged_db_collection(data)
 
+    if isinstance(data, dict) and isinstance(
+        next(iter(data.values()), None), new_sequence.Sequence
+    ):
+        # if we have a dict of Sequences, {name: seq, ...}, then the provided names
+        # may differ to the name attribute on the Sequences. If the Sequences have
+        # annotations, then we will need to map between the names in order to
+        # query the annotation database. We do this by creating a name_map.
+        # Note that we do (seq.name or name) to handle when the Sequence name is
+        # None
+        name_map = {name: (seq.name or name) for name, seq in data.items()}
+        data = {(seq.name or name): numpy.array(seq) for name, seq in data.items()}
+
     seqs_data = coerce_to_seqs_data_dict(data, label_to_name=label_to_name)
 
     moltype = new_moltype.get_moltype(moltype)
@@ -2329,6 +2344,7 @@ def make_unaligned_seqs(
         source=source,
         annotation_db=annotation_db,
         offset=offset,
+        name_map=name_map,
     )
 
 
@@ -2342,6 +2358,7 @@ def _(
     source: OptPathType = None,
     annotation_db: SupportsFeatures = None,
     offset: dict[str, int] = None,
+    name_map: dict[str, str] = None,
 ) -> SequenceCollection:
     moltype = new_moltype.get_moltype(moltype)
     if not moltype.is_compatible_alphabet(data.alphabet):
@@ -2361,6 +2378,7 @@ def _(
         info=info,
         annotation_db=annotation_db,
         source=source,
+        name_map=name_map,
     )
 
 
@@ -3556,7 +3574,7 @@ class Alignment(SequenceCollection):
             included
         """
         result = numpy.full((self.num_seqs, len(self)), False, dtype=bool)
-        for i, seqid in enumerate(self.names):
+        for i, seqid in enumerate(self._name_map.values()):
             if include_ambiguity:
                 seq_array = self._seqs_data.get_gapped_seq_array(seqid=seqid)
                 result[i][seq_array >= self.moltype.most_degen_alphabet().gap_index] = (
diff --git a/src/cogent3/parse/sequence.py b/src/cogent3/parse/sequence.py
index dbdfbc1c7..b77541a0e 100644
--- a/src/cogent3/parse/sequence.py
+++ b/src/cogent3/parse/sequence.py
@@ -142,3 +142,11 @@ def get_parser(fmt: str) -> typing.Callable[[SeqParserInputTypes], ParserOutputT
         return PARSERS[fmt]
     except KeyError:
         raise ValueError(f"Unsupported format {fmt!r}")
+
+
+def is_genbank(fmt: str) -> bool:
+    """whether the provided format is a genbank format"""
+    try:
+        return get_parser(fmt).__module__.endswith("genbank")
+    except ValueError:
+        return False


Test patch:
diff --git a/tests/test_core/test_core_standalone.py b/tests/test_core/test_core_standalone.py
index a25dc62a6..139124a79 100644
--- a/tests/test_core/test_core_standalone.py
+++ b/tests/test_core/test_core_standalone.py
@@ -923,3 +923,11 @@ def test_load_seq_moltype_offset(moltype, offset, DATA_DIR):
     )
     assert seq.moltype.label == moltype
     assert seq.annotation_offset == offset
+
+
+@pytest.mark.parametrize("ann_path", [True, False])
+@pytest.mark.parametrize("new_type", [True, False])
+def test_load_seq_genank(ann_path, new_type, DATA_DIR):
+    path = DATA_DIR / "annotated_seq.gb"
+    s = load_seq(path, annotation_path=path if ann_path else None, new_type=new_type)
+    assert len(s.annotation_db) == 10
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index ac52c2efc..fffaf3cdc 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -7,7 +7,7 @@ from warnings import catch_warnings, filterwarnings
 import numpy
 import pytest
 
-from cogent3 import get_app, load_unaligned_seqs, open_
+from cogent3 import get_app, load_seq, load_unaligned_seqs, open_
 from cogent3._version import __version__
 from cogent3.core import new_alignment, new_alphabet, new_moltype, new_sequence
 from cogent3.core.annotation import Feature
@@ -138,6 +138,16 @@ def gff_db(DATA_DIR):
     return load_annotations(path=DATA_DIR / "simple.gff")
 
 
+@pytest.fixture(scope="function")
+def annotated_seq(DATA_DIR):
+    return load_seq(
+        DATA_DIR / "c_elegans_WS199_dna_shortened.fasta",
+        annotation_path=DATA_DIR / "c_elegans_WS199_shortened_gff.gff3",
+        moltype="dna",
+        new_type=True,
+    )
+
+
 @pytest.fixture(scope="function")
 def seqcoll_db(DATA_DIR):
     from cogent3.parse.fasta import iter_fasta_records
@@ -1804,16 +1814,19 @@ def test_sequence_collection_dotplot(dotplot_seqs):
         dotplot_seqs.dotplot(name1="Human", name2="Dog")
 
 
-def test_sequence_collection_dotplot_annotated():
-    """exercising dotplot method with annotated sequences"""
-    db = GffAnnotationDb()
-    db.add_feature(seqid="Human", biotype="exon", name="fred", spans=[(10, 15)])
+@pytest.mark.parametrize("with_annotations", [True, False])
+def test_sequence_collection_dotplot_annotated(annotated_seq, with_annotations):
+    if not with_annotations:
+        annotated_seq.replace_annotation_db(None)  # this drops all annotations
 
-    data = {"Human": "CAGATTTGGCAGTT-", "Mouse": "CAGATTCAGCAGGTG"}
-    seqs = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    seqs.annotation_db = db
-    seqs = seqs.take_seqs(["Human", "Mouse"], copy_annotations=True)
-    _ = seqs.dotplot(show_progress=False)
+    coll = new_alignment.make_unaligned_seqs(
+        {"c_elegans": annotated_seq}, moltype="dna"
+    )
+    dp = coll.dotplot()
+    if with_annotations:
+        assert len(dp.figure.data) > 2
+    else:
+        assert len(dp.figure.data) == 2
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_parse/test_sequence.py b/tests/test_parse/test_sequence.py
index c95bf623f..a9cc56acd 100644
--- a/tests/test_parse/test_sequence.py
+++ b/tests/test_parse/test_sequence.py
@@ -64,3 +64,13 @@ def test_line_based_url(DATA_DIR):
     from_url = list(parser(url))
     from_local = list(parser(DATA_DIR / fname))
     assert from_url == from_local
+
+
+@pytest.mark.parametrize("fmt", ("gb", "gbk", "gbff", "genbank"))
+def test_is_genbank(fmt):
+    assert sequence.is_genbank(fmt)
+
+
+@pytest.mark.parametrize("fmt", ("blah", "fa", "xml", "nex", None))
+def test_is_not_genbank(fmt):
+    assert not sequence.is_genbank(fmt)