Repository: cogent3/cogent3
Pull request #2125 (2024-11-21 07:04:18)

ENH: implement eq and neq methods
## Summary by Sourcery

Implement equality and inequality methods for sequence-related classes to support object comparison and add corresponding tests to ensure their correct functionality.

Enhancements:
- Implement equality (__eq__) and inequality (__neq__) methods for various classes in the new_alignment and new_sequence modules to enable object comparison.

Tests:
- Add tests to verify the correct behavior of the newly implemented equality and inequality methods for sequence collections and alignments.


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 306747553..d31f1400c 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -4713,9 +4713,7 @@ class Alignment(AlignmentI, SequenceCollection):
         self.annotation_db.add_feature(**feature.to_dict())
         return result
 
-    def get_projected_features(
-        self, *, seqid: str, **kwargs
-    ) -> list[Feature]:  # ported
+    def get_projected_features(self, *, seqid: str, **kwargs) -> list[Feature]:
         """projects all features from other sequences onto seqid"""
         # todo gah should there be a generator version,
         # iter_projected_features()?
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 3afe4ebff..25c0db1e0 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -260,6 +260,12 @@ class SeqsDataABC(ABC):
         **kwargs,
     ): ...
 
+    @abstractmethod
+    def __eq__(self, value: object) -> bool: ...
+
+    @abstractmethod
+    def __ne__(self, value: object) -> bool: ...
+
     @abstractmethod
     def get_seq_length(self, seqid: str) -> int: ...
 
@@ -324,14 +330,12 @@ class SeqsData(SeqsDataABC):
         an instance of CharAlphabet valid for the sequences
     offset
         a dictionary of {name: offset} pairs indicating the offset of the sequence
-    reversed
-        a boolean indicating if the sequences are reversed
     check
         a boolean indicating if the data should be checked for naming consistency
         between arguments
     """
 
-    __slots__ = ("_data", "_alphabet", "_offset", "_reversed")
+    __slots__ = ("_data", "_alphabet", "_offset")
 
     def __init__(
         self,
@@ -357,6 +361,26 @@ class SeqsData(SeqsDataABC):
             arr.flags.writeable = False
             self._data[str(name)] = arr
 
+    def __eq__(self, other: SeqsDataABC) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        for attr_name in ("_alphabet", "_offset"):
+            self_attr = getattr(self, attr_name)
+            other_attr = getattr(other, attr_name)
+            if self_attr != other_attr:
+                return False
+
+        # compare individuals sequences
+        if self._data.keys() != other._data.keys():
+            return False
+        return all(
+            numpy.array_equal(self._data[name], other._data[name])
+            for name in self._data
+        )
+
+    def __ne__(self, other: object) -> bool:
+        return not self == other
+
     @classmethod
     def from_seqs(
         cls,
@@ -1990,10 +2014,16 @@ class SequenceCollection:
 
         return FORMATTERS["fasta"](self.to_dict())
 
-    def __eq__(
-        self, other: Union[SequenceCollection, dict]
-    ) -> bool:  # refactor: design
-        return id(self) == id(other)
+    def __eq__(self, other: SequenceCollection) -> bool:
+        self_init = self._get_init_kwargs()
+        other_init = other._get_init_kwargs()
+        for key, self_val in self_init.items():
+            if key in ("annotation_db", "slice_record"):
+                continue
+            other_val = other_init.get(key)
+            if self_val != other_val:
+                return False
+        return True
 
     def __ne__(self, other: SequenceCollection) -> bool:
         return not self.__eq__(other)
@@ -3132,6 +3162,27 @@ class AlignedSeqsData(AlignedSeqsDataABC):
             if len(names) != gapped_seqs.shape[0]:
                 raise ValueError(f"{len(names)=} != {gapped_seqs.shape[0]=}")
 
+    def __eq__(self, other: AlignedSeqsDataABC) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        attrs = (
+            "_names",
+            "_name_to_index",
+            "_alphabet",
+            "_align_len",
+            "_offset",
+        )
+        for attr_name in attrs:
+            self_attr = getattr(self, attr_name)
+            other_attr = getattr(other, attr_name)
+            if self_attr != other_attr:
+                return False
+
+        return numpy.all(self._gapped == other._gapped)
+
+    def __ne__(self, other: object) -> bool:
+        return not self == other
+
     @classmethod
     def from_seqs(
         cls,
@@ -3884,6 +3935,12 @@ class Alignment(SequenceCollection):
     def _post_init(self):
         self._seqs = _IndexableSeqs(self, make_seq=self._make_aligned)
 
+    def __eq__(self, other: "Alignment") -> bool:
+        return super().__eq__(other) and self._slice_record == other._slice_record
+
+    def __ne__(self, other: "Alignment") -> bool:
+        return not self == other
+
     def _get_init_kwargs(self) -> dict:
         """returns the kwargs needed to re-instantiate the object"""
         return {
@@ -5373,15 +5430,6 @@ class Alignment(SequenceCollection):
         self.annotation_db.add_feature(**feature.to_dict())
         return result
 
-    def get_projected_features(self, *, seqid: str, **kwargs) -> list[Feature]:
-        """projects all features from other sequences onto seqid"""
-        annots = []
-        for name in self.names:
-            if name == seqid:
-                continue
-            annots.extend(list(self.get_features(seqid=name, **kwargs)))
-        return [self.get_projected_feature(seqid=seqid, feature=a) for a in annots]
-
     def get_drawables(
         self, *, biotype: Optional[str, typing.Iterable[str]] = None
     ) -> dict:
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index 265743935..09832ad8a 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -2081,6 +2081,11 @@ class SliceRecordABC(ABC):
 
     __slots__ = ("start", "stop", "step", "_offset")
 
+    @abstractmethod
+    def __eq__(self, other): ...
+
+    @abstractmethod
+    def __neq__(self, other): ...
     @property
     @abstractmethod
     def parent_len(self) -> int: ...
@@ -2515,6 +2520,18 @@ class SliceRecord(SliceRecordABC):
         self._parent_len = parent_len
         self._offset = offset or 0
 
+    def __eq__(self, other: SliceRecordABC) -> bool:
+        return (
+            self.start == other.start
+            and self.stop == other.stop
+            and self.step == other.step
+            and self._parent_len == other.parent_len
+            and self._offset == other.offset
+        )
+
+    def __neq__(self, other: SliceRecordABC) -> bool:
+        return self != other
+
     @property
     def parent_len(self) -> int:
         return self._parent_len


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index fb4a313cf..9444555cc 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -5029,3 +5029,55 @@ def test_alignment_deepcopy(simple_aln):
     assert set(copied.names) == set(renamed.names)
     # and can get a sequence with a new name
     assert str(copied.seqs["A"]) == str(renamed.seqs["A"])
+
+
+@pytest.mark.parametrize(
+    "mk_cls",
+    [
+        new_alignment.make_unaligned_seqs,
+        new_alignment.make_aligned_seqs,
+    ],
+)
+def test_collections_equal(aligned_dict, mk_cls):
+    coll1 = mk_cls(aligned_dict, moltype="dna")
+    coll2 = mk_cls(aligned_dict, moltype="dna")
+    assert coll1 is not coll2
+    assert coll1 == coll2
+
+
+@pytest.mark.parametrize(
+    "mk_cls",
+    [
+        new_alignment.make_unaligned_seqs,
+        new_alignment.make_aligned_seqs,
+    ],
+)
+@pytest.mark.parametrize(
+    "kwargs",
+    [dict(moltype="protein"), dict(is_reversed=True), dict(name_map=dict(A="a")), {}],
+)
+def test_collections_not_equal(aligned_dict, mk_cls, kwargs):
+    coll1 = mk_cls(aligned_dict, moltype="dna")
+    aligned_dict = aligned_dict if kwargs else {**aligned_dict, **{"seq1": "TTTTTT"}}
+    kwargs = {**dict(moltype="dna"), **kwargs}
+    coll2 = mk_cls(aligned_dict, **kwargs)
+    assert coll1 != coll2
+
+
+def test_alignment_not_equal_sliced(aligned_dict):
+    coll1 = new_alignment.make_aligned_seqs(aligned_dict, moltype="dna")
+    coll2 = coll1[:2]
+    assert coll1 != coll2
+
+
+@pytest.mark.parametrize(
+    "type1,type2",
+    [
+        (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+        (new_alignment.make_aligned_seqs, new_alignment.make_unaligned_seqs),
+    ],
+)
+def test_alignment_not_equal_types(aligned_dict, type1, type2):
+    coll1 = type1(aligned_dict, moltype="dna")
+    coll2 = type2(aligned_dict, moltype="dna")
+    assert coll1 != coll2
diff --git a/tests/test_core/test_new_aln_annotation.py b/tests/test_core/test_new_aln_annotation.py
index 9cdd81b64..5b2fb53f1 100644
--- a/tests/test_core/test_new_aln_annotation.py
+++ b/tests/test_core/test_new_aln_annotation.py
@@ -502,30 +502,3 @@ def test_copy_annotations_no_db(gff_db, mk_cls):
 
     seq_coll.copy_annotations(gff_db)
     assert seq_coll.annotation_db.num_matches() == gff_db.num_matches()
-
-
-def test_project_features_onto_specified_seqid():
-    db = GffAnnotationDb()
-    aln = new_alignment.make_aligned_seqs(
-        {"x": "-AAAAAAAAA", "y": "TTTT--TTTT"}, moltype="dna"
-    )
-    db.add_feature(seqid="x", biotype="exon", name="fred", spans=[(3, 8)])
-    aln.annotation_db = db
-
-    exons = aln.get_projected_features(seqid="y", biotype="exon")
-    assert len(exons) == 1
-    assert str(aln.get_seq("y")[exons[0].map.without_gaps()]), "TTT"
-    assert "biotype='exon', name='fred', map=[-2-, 4:7]/8" in str(exons[0])
-
-
-def test_project_features_no_features_for_specified_seqid():
-    db = GffAnnotationDb()
-    aln = new_alignment.make_aligned_seqs(
-        {"seq1": "ATGCGT", "seq2": "AT--GT"}, moltype="dna"
-    )
-    db.add_feature(seqid="seq1", biotype="gene", name="gene1", spans=[(0, 6)])
-    aln.annotation_db = db
-
-    projected_features = aln.get_projected_features(seqid="seq2", biotype="exon")
-
-    assert len(projected_features) == 0