Repository: cogent3/cogent3
Pull request #2273 (2025-03-27 07:12:16)

BUG: fixing behaviour of masking annotations
## Summary by Sourcery

Fix and improve the behavior of masking annotations in sequences and alignments

Bug Fixes:
- Corrected the masking of annotations to handle cases with no matching features
- Improved handling of masking for specific sequences in an alignment

Enhancements:
- Added support for masking annotations with more flexible options
- Implemented seqid-specific masking in alignments

Tests:
- Added new test cases for masking annotations with different scenarios
- Introduced tests for masking with missing features
- Added parametrized tests for shadow and non-shadow masking


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 21590e4eb..6b906ec19 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -6606,7 +6606,8 @@ class Alignment(SequenceCollection):
         biotypes: PySeqStr,
         mask_char: str = "?",
         shadow: bool = False,
-    ):
+        seqid: str | None = None,
+    ) -> typing_extensions.Self:
         """returns an alignment with regions replaced by mask_char
 
         Parameters
@@ -6628,9 +6629,12 @@ class Alignment(SequenceCollection):
         for seq in self.seqs:
             parent_name = self._name_map[seq.name]
             gaps[parent_name] = seq.map.array
-            ungapped[parent_name] = numpy.array(
-                seq.seq.with_masked_annotations(biotypes, mask_char, shadow),
-            )
+            if seqid is None or seq.name == seqid:
+                ungapped[parent_name] = numpy.array(
+                    seq.seq.with_masked_annotations(biotypes, mask_char, shadow),
+                )
+            else:
+                ungapped[parent_name] = numpy.array(seq.seq)
 
         seq_data = self._seqs_data.from_seqs_and_gaps(
             seqs=ungapped,
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index bd2e0832a..a58108fd8 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -1363,12 +1363,17 @@ class Sequence:
             annotation_db=None if exclude_annotations else self.annotation_db,
         )
 
+    @c3warn.deprecated_args(
+        "2025.6",
+        old_new=[("annot_types", "biotypes")],
+        reason="consistency with Alignment method",
+    )
+    @c3warn.deprecated_args("2025.6", reason="not used", discontinued="extend_query")
     def with_masked_annotations(
         self,
-        annot_types: StrORIterableStr,
+        biotypes: StrORIterableStr,
         mask_char: str | None = None,
         shadow: bool = False,
-        extend_query: bool = False,
     ) -> typing_extensions.Self:
         """returns a sequence with annot_types regions replaced by mask_char
         if shadow is False, otherwise all other regions are masked.
@@ -1394,23 +1399,15 @@ class Sequence:
         )
 
         annotations = []
-        annot_types = [annot_types] if isinstance(annot_types, str) else annot_types
-        for annot_type in annot_types:
+        biotypes = [biotypes] if isinstance(biotypes, str) else biotypes
+        for annot_type in biotypes:
             annotations += list(
                 self.get_features(biotype=annot_type, allow_partial=True),
             )
-
         if not annotations:
-            region = Feature(
-                parent=self,
-                seqid=self.name,
-                name=None,
-                biotype=None,
-                map=FeatureMap.from_locations(locations=[], parent_length=len(self)),
-                strand="+",
-            )
-        else:
-            region = annotations[0].union(annotations[1:])
+            return self
+
+        region = annotations[0].union(annotations[1:])
 
         if shadow:
             region = region.shadow()


Test patch:
diff --git a/tests/test_core/test_new_aln_annotation.py b/tests/test_core/test_new_aln_annotation.py
index 0d21f9e38..77519d961 100644
--- a/tests/test_core/test_new_aln_annotation.py
+++ b/tests/test_core/test_new_aln_annotation.py
@@ -45,7 +45,7 @@ def makeSampleAlignment():
 
 
 @pytest.fixture
-def ann_aln():
+def ann_aln1():
     # synthetic annotated alignment
     return makeSampleAlignment()
 
@@ -214,8 +214,9 @@ def test_seqcoll_query(seqcoll_db):
 
 def test_feature_projection_ungapped(ann_aln):
     # projection onto ungapped sequence
+    ann_aln = makeSampleAlignment()
     expecteds = {"FAKE01": "CCCAAAATTTTTT", "FAKE02": "CCC-----TTTTT"}
-    aln_ltr = next(iter(ann_aln.get_features(biotype="LTR")))
+    aln_ltr = next(iter(ann_aln.get_features(biotype="LTR", on_alignment=True)))
     seq_name = "FAKE01"
     expected = expecteds[seq_name]
     num = ann_aln.annotation_db.num_matches()
@@ -227,13 +228,14 @@ def test_feature_projection_ungapped(ann_aln):
     assert seq_ltr.parent == ann_aln.get_seq(seq_name)
 
 
-def test_feature_projection_gapped(ann_aln):
+def test_feature_projection_gapped(ann_aln1):
+    aln = ann_aln1
     # projection onto gapped sequence
     expecteds = {"FAKE01": "CCCAAAATTTTTT", "FAKE02": "CCC-----TTTTT"}
-    aln_ltr = next(iter(ann_aln.get_features(biotype="LTR")))
+    aln_ltr = next(iter(aln.get_features(biotype="LTR")))
     seq_name = "FAKE02"
     expected = expecteds[seq_name]
-    seq_ltr = ann_aln.get_projected_feature(seqid=seq_name, feature=aln_ltr)
+    seq_ltr = aln.get_projected_feature(seqid=seq_name, feature=aln_ltr)
 
     with pytest.raises(ValueError):
         seq_ltr.get_slice(complete=True)
@@ -243,7 +245,7 @@ def test_feature_projection_gapped(ann_aln):
     expected = expected.replace("-", "")
     assert str(seq_ltr.get_slice()) == expected
     assert seq_ltr.seqid == seq_name
-    assert seq_ltr.parent == ann_aln.get_seq(seq_name)
+    assert seq_ltr.parent == aln.get_seq(seq_name)
 
 
 def test_get_feature():
@@ -568,17 +570,8 @@ def test_project_features_no_features_for_specified_seqid():
     assert len(projected_features) == 0
 
 
-def test_annotated_region_masks():
-    """masking a sequence with specific features"""
-    # ported from test_features.py
-    # refactor: break into multiple tests
-
-    # Annotated regions can be masked (observed sequence characters
-    # replaced by another), either through the sequence on which they
-    # reside or by projection from the alignment. Note that mask_char must
-    # be a valid character for the sequence MolType. Either the features
-    # (multiple can be named), or their shadow, can be masked.
-
+@pytest.fixture
+def ann_aln():
     # We create an alignment with a sequence that has two different annotation types.
     orig_data = {"x": "C-CCCAAAAAGGGAA", "y": "-T----TTTTG-GTT"}
     db = GffAnnotationDb()
@@ -590,10 +583,21 @@ def test_annotated_region_masks():
         seqid="x",
     )
     db.add_feature(seqid="y", biotype="repeat", name="frog", spans=[(5, 7)])
-    aln = new_alignment.make_aligned_seqs(orig_data, moltype="dna")
-    aln.annotation_db = db
+    return new_alignment.make_aligned_seqs(orig_data, moltype="dna", annotation_db=db)
+
+
+def test_annotated_region_masks(ann_aln):
+    """masking a sequence with specific features"""
+    # ported from test_features.py
+    # refactor: break into multiple tests
 
-    assert aln.to_dict() == {"x": "C-CCCAAAAAGGGAA", "y": "-T----TTTTG-GTT"}
+    # Annotated regions can be masked (observed sequence characters
+    # replaced by another), either through the sequence on which they
+    # reside or by projection from the alignment. Note that mask_char must
+    # be a valid character for the sequence MolType. Either the features
+    # (multiple can be named), or their shadow, can be masked.
+    raw_data = ann_aln.to_dict()
+    aln = ann_aln
     x = aln.get_seq("x")
     y = aln.get_seq("y")
     exon = next(iter(x.get_features(biotype="exon")))
@@ -654,33 +658,54 @@ def test_annotated_region_masks():
     )
 
     # The same methods can be applied to annotated Alignment's.
-
-    assert aln.with_masked_annotations("exon", mask_char="?").to_dict() == {
-        "x": "?-???AAAAAGGGAA",
-        "y": "-T----TTTTG-GTT",
-    }
-    assert aln.with_masked_annotations(
-        "exon",
-        mask_char="?",
-        shadow=True,
-    ).to_dict() == {"x": "C-CCC??????????", "y": "-?----?????-???"}
-    assert aln.with_masked_annotations("repeat", mask_char="?").to_dict() == {
-        "x": "C-CCCAAAAA???AA",
-        "y": "-T----TTTT?-?TT",
-    }
-    assert aln.with_masked_annotations(
+    # only x has an exon, both have repeats
+    masked = aln.with_masked_annotations("exon", mask_char="?")
+    assert masked.to_dict() == {"x": "?-???AAAAAGGGAA", "y": raw_data["y"]}
+    masked = aln.with_masked_annotations("exon", mask_char="?", shadow=True)
+    assert masked.to_dict() == {"x": "C-CCC??????????", "y": raw_data["y"]}
+    masked = aln.with_masked_annotations("repeat", mask_char="?")
+    assert masked.to_dict() == {"x": "C-CCCAAAAA???AA", "y": "-T----TTTT?-?TT"}
+    masked = aln.with_masked_annotations(
         "repeat",
         mask_char="?",
         shadow=True,
-    ).to_dict() == {"x": "?-????????GGG??", "y": "-?----????G-G??"}
-    assert aln.with_masked_annotations(["repeat", "exon"], mask_char="?").to_dict() == {
-        "x": "?-???AAAAA???AA",
-        "y": "-T----TTTT?-?TT",
-    }
-    assert aln.with_masked_annotations(["repeat", "exon"], shadow=True).to_dict() == {
-        "x": "C-CCC?????GGG??",
-        "y": "-?----????G-G??",
+    )
+    assert masked.to_dict() == {"x": "?-????????GGG??", "y": "-?----????G-G??"}
+    masked = aln.with_masked_annotations(["repeat", "exon"], mask_char="?")
+    assert masked.to_dict() == {"x": "?-???AAAAA???AA", "y": "-T----TTTT?-?TT"}
+    masked = aln.with_masked_annotations(["repeat", "exon"], shadow=True)
+    assert masked.to_dict() == {"x": "C-CCC?????GGG??", "y": "-?----????G-G??"}
+
+
+def test_with_masked_one_seqid():
+    raw_data = {
+        "x": "AACCCAAAATTTTTTGGGGGGGGGGCCCC",
+        "y": "AACCC-----TTTTTGGGGGGGGGGCC--",
     }
+    aln = new_alignment.make_aligned_seqs(raw_data, moltype="dna")
+    start, stop = 2, 10
+    aln.annotation_db.add_feature(
+        biotype="repeat",
+        name="blah",
+        spans=[(start, stop)],
+        seqid="y",
+    )
+    masked = aln.with_masked_annotations(biotypes="repeat", mask_char="?", seqid="y")
+    expect = {"x": raw_data["x"], "y": "AA???-----?????GGGGGGGGGGCC--"}
+    assert masked.to_dict() == expect
+
+
+@pytest.mark.parametrize("shadow", [True, False])
+def test_with_masked_missing_feature(ann_aln1, shadow):
+    aln = ann_aln1
+    expect = aln.to_dict()
+    # with no matches, should return the original sequences
+    masked = aln.with_masked_annotations(
+        biotypes="not-present",
+        mask_char="?",
+        shadow=shadow,
+    )
+    assert masked.to_dict() == expect
 
 
 def test_masking_strand_agnostic_aln():
diff --git a/tests/test_core/test_new_seq_annotation.py b/tests/test_core/test_new_seq_annotation.py
index 105b9b08c..f95a51d28 100644
--- a/tests/test_core/test_new_seq_annotation.py
+++ b/tests/test_core/test_new_seq_annotation.py
@@ -180,3 +180,25 @@ def test_feature_names_seq():
     s = seq[f]
     assert s.name == "s1-cds"
     assert s.name == f.name
+
+
+def test_seq_with_masked_annotations():
+    seq = makeSampleSequence("seq1", with_gaps=False)
+    raw_seq = str(seq)
+    cds = next(iter(seq.annotation_db.get_records_matching(biotype="CDS")))
+    start, stop = cds["start"], cds["stop"]
+    expect = raw_seq[:start] + "?" * (stop - start) + raw_seq[stop:]
+    masked = seq.with_masked_annotations(biotypes="CDS")
+    assert str(masked) == expect
+    shadow = seq.with_masked_annotations(biotypes="CDS", shadow=True)
+    expect = "?" * start + raw_seq[start:stop] + "?" * (len(raw_seq) - stop)
+    assert str(shadow) == expect
+    print(shadow)
+
+
+@pytest.mark.parametrize("shadow", [True, False])
+def test_seq_with_masked_annotations_missing_feature(shadow):
+    seq = makeSampleSequence("seq1", with_gaps=False)
+    raw_seq = str(seq)
+    masked = seq.with_masked_annotations(biotypes="not-present", shadow=shadow)
+    assert str(masked) == raw_seq