Repository: cogent3/cogent3
Pull request #2241 (2025-03-14 03:40:16)

ENH: support storing extra attributes on Features
## Summary by Sourcery

New Features:
- Adds the `xattr` property to the `Feature` class, allowing developers to store additional data without modifying the core API.


Patch:
diff --git a/changelog.d/20250314_131212_Gavin.Huttley.md b/changelog.d/20250314_131212_Gavin.Huttley.md
index 21647c6d5..9c674c666 100644
--- a/changelog.d/20250314_131212_Gavin.Huttley.md
+++ b/changelog.d/20250314_131212_Gavin.Huttley.md
@@ -9,12 +9,14 @@ Uncomment the section that is right (remove the HTML comment wrapper).
 - First time contributions from @c-jonesy, @Yixuan567, @Qinzi27
   and @je634 👏🚀🎉
 
-<!--
+
 ### ENH
 
-- A bullet item for the ENH category.
+- The cogent3.core.annotation.Feature class now gets an xattr property.
+  This allows developers to pack additional data into the Feature object
+  without compromising the core API.
+
 
--->
 <!--
 ### BUG
 
diff --git a/src/cogent3/core/annotation.py b/src/cogent3/core/annotation.py
index 032a92eae..e5da65503 100644
--- a/src/cogent3/core/annotation.py
+++ b/src/cogent3/core/annotation.py
@@ -1,9 +1,18 @@
+import typing
 from collections.abc import Iterable
 
+import typing_extensions
 from numpy import array
 
 from .location import FeatureMap
 
+if typing.TYPE_CHECKING:
+    from cogent3.core.new_alignment import Alignment
+    from cogent3.core.new_sequence import Sequence
+    from cogent3.draw.drawable import Shape
+
+SeqORAlign = typing.Union["Sequence", "Alignment"]
+
 
 # TODO gah write docstrings!
 class Feature:
@@ -19,18 +28,20 @@ class Feature:
         "_seqid",
         "_serialisable",
         "_strand",
+        "_xattr",
     )
 
     # TODO gah implement a __new__ to trap args for serialisation purposes?
     def __init__(
         self,
         *,
-        parent,
+        parent: SeqORAlign,
         seqid: str,
         map: FeatureMap,
         biotype: str,
         name: str,
         strand: str,
+        xattr: dict[str, typing.Any] | None = None,
     ) -> None:
         # _serialisable is used for creating derivative instances
         d = locals()
@@ -45,32 +56,33 @@ class Feature:
         data.extend((self.seqid, self.biotype, self.name))
         self._id = hash(tuple(data))
         self._strand = strand
+        self._xattr = xattr
 
-    def __eq__(self, other):
+    def __eq__(self, other: typing_extensions.Self) -> bool:
         return self._id == other._id
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         """Features can be used in a dictionary!"""
         return self._id
 
     @property
-    def parent(self):
+    def parent(self) -> SeqORAlign:
         return self._parent
 
     @property
-    def seqid(self):
+    def seqid(self) -> str:
         return self._seqid
 
     @property
-    def map(self):
+    def map(self) -> FeatureMap:
         return self._map
 
     @property
-    def biotype(self):
+    def biotype(self) -> str:
         return self._biotype
 
     @property
-    def name(self):
+    def name(self) -> str:
         return self._name
 
     def get_slice(
@@ -78,7 +90,7 @@ class Feature:
         complete: bool = False,
         allow_gaps: bool = False,
         apply_name: bool = True,
-    ):
+    ) -> SeqORAlign:
         """
         The corresponding sequence fragment.
 
@@ -112,7 +124,7 @@ class Feature:
         result = self.parent[fmap.start : fmap.end]
         return self._do_seq_slice(result, apply_name)
 
-    def _do_seq_slice(self, result, apply_name):
+    def _do_seq_slice(self, result: SeqORAlign, apply_name: bool) -> SeqORAlign:
         if self.reversed:
             result = result.rc()
         if self.map.num_spans > 1:
@@ -122,7 +134,7 @@ class Feature:
             result.name = self.name
         return result
 
-    def without_lost_spans(self):
+    def without_lost_spans(self) -> typing_extensions.Self:
         """Keeps only the parts which are actually present in the underlying sequence"""
         if self.map.complete:
             return self
@@ -130,7 +142,7 @@ class Feature:
         kwargs = {**self._serialisable, "map": self.map[keep]}
         return self.__class__(**kwargs)
 
-    def as_one_span(self):
+    def as_one_span(self) -> typing_extensions.Self:
         """returns a feature that preserves any gaps"""
         kwargs = {
             **self._serialisable,
@@ -139,7 +151,7 @@ class Feature:
         }
         return self.__class__(**kwargs)
 
-    def shadow(self):
+    def shadow(self) -> typing_extensions.Self:
         """returns new instance corresponding to disjoint of self coordinates"""
         kwargs = {
             **self._serialisable,
@@ -175,12 +187,16 @@ class Feature:
         }
         return self.__class__(**kwargs)
 
-    def get_coordinates(self):
+    def get_coordinates(self) -> list[tuple[int, int]]:
         """returns sequence coordinates of this Feature as
         [(start1, end1), ...]"""
         return self.map.get_coordinates()
 
-    def get_children(self, biotype: str | None = None, **kwargs):
+    def get_children(
+        self,
+        biotype: str | None = None,
+        **kwargs,
+    ) -> typing.Iterable[typing_extensions.Self]:
         """generator returns sub-features of self optionally matching biotype"""
         offset = getattr(self.parent, "annotation_offset", 0)
         start = self.map.start + offset
@@ -204,7 +220,7 @@ class Feature:
                 continue
             yield feature
 
-    def get_parent(self, **kwargs):
+    def get_parent(self, **kwargs) -> typing.Iterable[typing_extensions.Self]:
         """generator returns parent features of self optionally matching biotype"""
         offset = getattr(self.parent, "annotation_offset", 0)
         start = self.map.start + offset
@@ -227,7 +243,7 @@ class Feature:
                 continue
             yield feature
 
-    def union(self, features: Iterable):
+    def union(self, features: Iterable) -> typing_extensions.Self:
         """return as a single Feature
 
         Notes
@@ -278,23 +294,32 @@ class Feature:
 
         return self.__class__(**kwargs)
 
-    def get_drawable(self):
+    def get_drawable(self) -> "Shape":
         """returns plotly trace"""
         from cogent3.draw.drawable import make_shape
 
         return make_shape(type_=self, parent_length=len(self.parent))
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, typing.Any]:
         """returns"""
         result = {
             **self._serialisable,
             "spans": self.map.get_coordinates(),
         }
-        for key in ("map", "parent"):
+        for key in ("map", "parent", "xattr"):
             result.pop(key, None)
         return result
 
     @property
-    def reversed(self):
+    def reversed(self) -> bool:
         """whether Feature is on the reverse strand relative to bound object"""
         return self._strand == "-"
+
+    @property
+    def xattr(self) -> dict[str, typing.Any] | None:
+        return self._xattr
+
+    @xattr.setter
+    def xattr(self, val: typing.Any) -> None:
+        msg = "setting xattr is not supported"
+        raise TypeError(msg)
diff --git a/src/cogent3/parse/ebi.py b/src/cogent3/parse/ebi.py
index e536c2838..f9204458f 100644
--- a/src/cogent3/parse/ebi.py
+++ b/src/cogent3/parse/ebi.py
@@ -1,4 +1,5 @@
 """Provide a parser for SwissProt EBI format files."""
+
 from cogent3.core.sequence import Sequence
 from cogent3.parse.record import FieldError, RecordError
 from cogent3.parse.record_finder import (
@@ -1559,5 +1560,3 @@ def EbiParser(
                 continue
 
         yield sequence, header
-
-


Test patch:
diff --git a/tests/test_core/test_features.py b/tests/test_core/test_features.py
index 7c4380028..04739853d 100755
--- a/tests/test_core/test_features.py
+++ b/tests/test_core/test_features.py
@@ -10,8 +10,8 @@ from cogent3.core.annotation_db import BasicAnnotationDb, GffAnnotationDb
 # Complete version of manipulating sequence annotations
 from cogent3.util.deserialise import deserialise_object
 
-ASCII = cogent3.get_moltype("text")
-DNA = cogent3.get_moltype("dna")
+ASCII = cogent3.get_moltype("text", new_type=True)
+DNA = cogent3.get_moltype("dna", new_type=True)
 
 
 class FeaturesTest(TestCase):
@@ -115,8 +115,8 @@ def test_copy_annotations():
     """copying features from a db"""
     aln = cogent3.make_aligned_seqs(
         data=[["x", "-AAAAAAAAA"], ["y", "TTTT--CCCT"]],
-        array_align=False,
         moltype="dna",
+        new_type=True,
     )
     db = GffAnnotationDb()
     db.add_feature(seqid="y", biotype="exon", name="A", spans=[(5, 8)])
@@ -129,8 +129,8 @@ def test_copy_annotations_onto_seq():
     """copying features onto a sequence"""
     aln = cogent3.make_aligned_seqs(
         data=[["x", "-AAAAAAAAA"], ["y", "TTTT--CCCT"]],
-        array_align=False,
         moltype="dna",
+        new_type=True,
     )
     db = BasicAnnotationDb()
     db.add_feature(seqid="y", biotype="exon", name="A", spans=[(5, 8)])
@@ -148,7 +148,7 @@ def test_feature_residue():
     aln = cogent3.make_aligned_seqs(
         data=[["x", "C-CCCAAAAA"], ["y", "-T----TTTT"]],
         moltype=DNA,
-        array_align=False,
+        new_type=True,
     )
     db = aln.annotation_db
     assert str(aln), ">x\nC-CCCAAAAA\n>y\n-T----TTTT\n"
@@ -247,8 +247,8 @@ def test_feature_query_parent_seq():
 def test_feature_query_child_aln():
     aln = cogent3.make_aligned_seqs(
         data=[["x", "-AAAGGGGGAAC-CT"], ["y", "TTTT--TTTTAGGGA"]],
-        array_align=False,
         moltype="dna",
+        new_type=True,
     )
     aln = _add_features(aln, on_alignment=True)
     gene = next(iter(aln.get_features(biotype="CDS")))
@@ -504,22 +504,6 @@ def test_nested_get_slice():
     assert str(s[f]) == str(s[12:17])
 
 
-def test_roundtrip_annotated_seq():
-    """should work for a seq that has been reverse complemented"""
-    # the key that exposed the bug was a gap in the middle of the sequence
-    seq = DNA.make_seq(
-        seq="AAAGGGGGAACCT",
-        name="x",
-    )
-    seq.add_feature(biotype="exon", name="E1", spans=[(3, 8)])
-    seq.add_feature(biotype="exon", name="E2", spans=[(10, 13)])
-
-    rseq = deserialise_object(seq.to_json())
-    orig_annots = {a.name: str(a.get_slice()) for a in seq.get_features()}
-    got_annots = {a.name: str(a.get_slice()) for a in rseq.get_features()}
-    assert got_annots == orig_annots
-
-
 def test_roundtrip_rc_annotated_align():
     """should work for an alignment that has been reverse complemented"""
     # the key that exposed the bug was a gap in the middle of the sequence
@@ -593,90 +577,6 @@ def test_masking_strand_agnostic_aln():
     }
 
 
-def test_roundtrip_json():
-    """features can roundtrip from json"""
-
-    seq = DNA.make_seq(seq="AAAAATATTATTGGGT")
-    seq.add_feature(biotype="exon", name="myname", spans=[(0, 5)])
-    got = seq.to_json()
-    new = deserialise_object(got)
-    feat = next(iter(new.get_features(biotype="exon")))
-    assert str(feat.get_slice()) == "AAAAA"
-
-    # now with a list span
-    seq = seq[3:]
-    got = seq.to_json()
-    new = deserialise_object(got)
-    assert new.annotation_offset == 3
-    feat = next(iter(new.get_features(biotype="exon", allow_partial=True)))
-    assert str(feat.get_slice()) == "AA"
-
-
-def test_roundtripped_alignment():
-    """Alignment with annotations roundtrips correctly"""
-    # annotations just on member sequences
-    aln = cogent3.make_aligned_seqs(
-        data=[["x", "-AAAAAAAAA"], ["y", "TTTT--TTTT"]],
-        array_align=False,
-        moltype="dna",
-    )
-    db = GffAnnotationDb()
-    db.add_feature(seqid="x", biotype="exon", name="fred", spans=[(3, 8)])
-    aln.annotation_db = db
-    seq_exon = next(iter(aln.get_features(seqid="x", biotype="exon")))
-    expect = seq_exon.get_slice()
-
-    json = aln.to_json()
-    new = deserialise_object(json)
-    got_exons = next(iter(new.get_features(seqid="x", biotype="exon")))
-    assert got_exons.get_slice().to_dict() == expect.to_dict()
-
-    # annotations just on alignment
-    aln = cogent3.make_aligned_seqs(
-        data=[["x", "-AAAAAGGGG"], ["y", "TTTT--CCCC"]],
-        array_align=False,
-        moltype="dna",
-    )
-
-    f = aln.add_feature(biotype="generic", name="no name", spans=[(1, 4), (6, 10)])
-    expect = f.get_slice().to_dict()
-    json = aln.to_json()
-    new = deserialise_object(json)
-    got = next(iter(new.get_features(biotype="generic")))
-    assert got.get_slice().to_dict() == expect
-    # annotations on both alignment and sequence
-    aln = cogent3.make_aligned_seqs(
-        data=[["x", "-AAAAAGGGG"], ["y", "TTTT--CCCC"]],
-        array_align=False,
-        moltype="dna",
-    )
-    db = GffAnnotationDb()
-    db.add_feature(
-        seqid=None,
-        biotype="generic",
-        name="no name",
-        spans=[(1, 4), (6, 10)],
-        on_alignment=True,
-    )
-    db.add_feature(seqid="x", biotype="exon", name="1", spans=[(3, 8)])
-    aln.annotation_db = db
-    json = aln.to_json()
-    new = deserialise_object(json)
-    ## get back the exon
-    seq_exon = next(iter(aln.get_features(seqid="x", biotype="exon")))
-    expect = seq_exon.get_slice().to_dict()
-    got_exons = next(iter(new.get_features(seqid="x", biotype="exon")))
-    assert got_exons.get_slice().to_dict() == expect
-    ## get back the generic
-    expect = f.get_slice().to_dict()
-    got = next(iter(new.get_features(biotype="generic")))
-    assert got.get_slice().to_dict() == expect
-
-    # check masking of seq features still works
-    new = new.with_masked_annotations("exon", mask_char="?")
-    assert new[4:9].to_dict() == {"x": "?????", "y": "--CCC"}
-
-
 def test_feature_out_range():
     """features no longer included in an alignment will not be returned"""
     aln = cogent3.make_aligned_seqs(
@@ -710,8 +610,8 @@ def test_roundtripped_alignment_with_slices():
     # annotations just on member sequences
     aln = cogent3.make_aligned_seqs(
         data=[["x", "-AAAGGGGGAACCCT"], ["y", "TTTT--TTTTAGGGA"]],
-        array_align=False,
         moltype="dna",
+        new_type=True,
     )
     db = GffAnnotationDb()
     db.add_feature(seqid="x", biotype="exon", name="E1", spans=[(3, 8)])
@@ -723,12 +623,10 @@ def test_roundtripped_alignment_with_slices():
     sub_aln = aln[:-3]
     feats = list(sub_aln.get_features(biotype="exon", allow_partial=True))
     assert len(feats) == 2
+    # new type alignments DO NOT support serialising annotation_db's
     new = deserialise_object(sub_aln.to_json())
     feats = list(new.get_features(biotype="exon", allow_partial=True))
-    assert len(feats) == 2
-    gf1, gf2 = feats
-    assert gf1.get_slice().to_dict() == {"x": "GGGGG", "y": "--TTT"}
-    assert gf2.get_slice().to_dict() == {"x": "C", "y": "G"}
+    assert not len(feats)
 
 
 def test_feature_reverse():
@@ -742,7 +640,10 @@ def test_feature_reverse():
     # and show that after getting the reverse complement we have
     # exactly the same result from getting the CDS annotation.
 
-    plus = DNA.make_seq(seq="AAGGGGAAAACCCCCAAAAAAAAAATTTTTTTTTTAAA", name="plus")
+    plus = DNA.make_seq(
+        seq="AAGGGGAAAACCCCCAAAAAAAAAATTTTTTTTTTAAA",
+        name="plus",
+    )
     plus_cds = plus.add_feature(
         biotype="CDS",
         name="gene",
@@ -756,7 +657,7 @@ def test_feature_reverse():
 
 @pytest.mark.parametrize("moltype", ["protein", "bytes", "text"])
 def test_rc_feature_on_wrong_moltype(moltype):
-    moltype = cogent3.get_moltype(moltype)
+    moltype = cogent3.get_moltype(moltype, new_type=True)
     seq = moltype.make_seq(seq="AAGGGGAAAACCCCCAAAAAAAAAATTTTTTTTTTAAA", name="s1")
     cds = seq.add_feature(
         biotype="CDS",
@@ -822,7 +723,11 @@ def test_seq_degap_preserves_annotations():
 
 @pytest.mark.parametrize("aligned", [True, False])
 def test_align_degap_preserves_annotations(aligned):
-    kwargs = {"data": {"seq1": "GATN--", "seq2": "?GATCT"}, "moltype": DNA}
+    kwargs = {
+        "data": {"seq1": "GATN--", "seq2": "?GATCT"},
+        "moltype": DNA,
+        "new_type": True,
+    }
     coll = (
         cogent3.make_aligned_seqs(array_align=aligned, **kwargs)
         if aligned
@@ -834,3 +739,45 @@ def test_align_degap_preserves_annotations(aligned):
     got = coll.degap()
     assert got.annotation_db is coll.annotation_db
     assert len(got.annotation_db) == 1
+
+
+@pytest.fixture
+def feature_data():
+    from cogent3.core.location import FeatureMap
+
+    seq = DNA.make_seq(seq="ACGGTG", name="demo")
+    fmap = FeatureMap.from_locations(locations=[(0, 2), (4, 6)], parent_length=6)
+    return {
+        "seqid": "1",
+        "strand": 0,
+        "map": fmap,
+        "biotype": "repeat",
+        "name": "trf",
+        "parent": seq,
+    }
+
+
+@pytest.fixture
+def feature_data_xattr(feature_data):
+    feature_data["xattr"] = {
+        "repeat_type": "Tandem repeats",
+        "repeat_class": "trf",
+        "repeat_name": "trf",
+    }
+    return feature_data
+
+
+def test_feature_xattr(feature_data_xattr):
+    feature = Feature(**feature_data_xattr)
+    assert feature.xattr == feature_data_xattr["xattr"]
+
+
+def test_feature_xttar_none(feature_data):
+    feature = Feature(**feature_data)
+    assert feature.xattr is None
+
+
+def test_feature_xattr_set(feature_data):
+    feature = Feature(**feature_data)
+    with pytest.raises(TypeError):
+        feature.xattr = {}