Repository: cogent3/cogent3
Pull request #1838 (2024-04-30 03:16:44)

Add `__repr__` and `_repr_html_` to `SequenceCollection`
Adds `__repr__` and `_repr_html_` to `SequenceCollection` resolving #1682.

`__repr__` has been tested. `_repr_html_` has only been tested through the existing (slightly modified) `test_set_wrap_affects_repr_html` tests. Further explicit testing may be required in `SequenceCollectionTests`


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 7ee06c8e8..14d51588e 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -32,7 +32,16 @@ from collections import Counter, defaultdict
 from copy import deepcopy
 from functools import total_ordering
 from itertools import combinations
-from typing import Any, Iterable, Iterator, List, Optional, Tuple, Union
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import numpy
 
@@ -708,7 +717,7 @@ class _SequenceCollectionBase:
 
     # access as attribute if using default order.
 
-    def take_seqs(self, seqs, negate=False, **kwargs):
+    def take_seqs(self, seqs: Union[str, typing.Sequence[str]], negate=False, **kwargs):
         """Returns new Alignment containing only specified seqs.
 
         Note that the seqs in the new alignment will be references to the
@@ -2145,6 +2154,213 @@ class SequenceCollection(_SequenceCollectionBase):
 
         return dist_calc_app(self)
 
+    def __repr__(self):
+        if len(self.names) == 0:
+            return f"0x () {self.moltype.get_type()} seqcollection"
+
+        seqs = []
+        limit = 10
+        delimiter = ""
+
+        repr_seq_names = [min(self.names, key=lambda name: len(self.named_seqs[name]))]
+        if len(self.names) > 1:
+            # In case of a tie, min and max return first.
+            # reversed ensures if all seqs are of same length, different seqs are returned
+            repr_seq_names.append(
+                max(reversed(self.names), key=lambda name: len(self.named_seqs[name]))
+            )
+
+        for name in repr_seq_names:
+            elts = list(str(self.named_seqs[name])[: limit + 1])
+            if len(elts) > limit:
+                elts.append("...")
+            seqs.append(f"{name}[{delimiter.join(elts)}]")
+
+        if len(self.names) > 2:
+            seqs.insert(1, "...")
+
+        seqs = ", ".join(seqs)
+
+        return f"{len(self.names)}x ({seqs}) {self.moltype.get_type()} seqcollection"
+
+    def _repr_html_(self) -> str:
+        settings = self._repr_policy.copy()
+        env_vals = get_setting_from_environ(
+            "COGENT3_ALIGNMENT_REPR_POLICY",
+            dict(num_seqs=int, num_pos=int, wrap=int),
+        )
+        settings.update(env_vals)
+        return self.to_html(
+            name_order=self.names[: settings["num_seqs"]],
+            limit=settings["num_pos"],
+            wrap=settings["wrap"],
+        )
+
+    def to_html(
+        self,
+        name_order: Optional[typing.Sequence[str]] = None,
+        wrap: int = 60,
+        limit: Optional[int] = None,
+        colors: Optional[Mapping[str, str]] = None,
+        font_size: int = 12,
+        font_family: str = "Lucida Console",
+    ) -> str:
+        """returns html with embedded styles for sequence colouring
+
+        Parameters
+        ----------
+        name_order
+            order of names for display.
+        wrap
+            number of alignment columns per row
+        limit
+            truncate alignment to this length
+        colors
+            {character
+            moltype.
+        font_size
+            in points. Affects labels and sequence and line spacing
+            (proportional to value)
+        font_family
+            string denoting font family
+
+        Examples
+        ---------
+
+        In a jupyter notebook, this code is used to provide the representation.
+
+        .. code-block:: python
+
+            seq_col # is rendered by jupyter
+
+        You can directly use the result for display in a notebook as
+
+        .. code-block:: python
+
+            from IPython.core.display import HTML
+            HTML(seq_col.to_html())
+        """
+        css, styles = self.moltype.get_css_style(
+            colors=colors, font_size=font_size, font_family=font_family
+        )
+
+        # Gather stats about the sequence lengths
+        seq_lens = sorted([len(seq) for seq in self.seqs])
+        if seq_lens:
+            min_seq_len = seq_lens[0]
+            max_seq_len = seq_lens[-1]
+
+            # Find median
+            if len(seq_lens) % 2 == 0:
+                total = (
+                    seq_lens[(len(seq_lens) - 2) // 2] + seq_lens[len(seq_lens) // 2]
+                )
+                med_seq_len = total // 2 if total % 2 == 0 else round(total / 2, 1)
+            else:
+                med_seq_len = seq_lens[(len(seq_lens) - 1) // 2]
+        else:
+            min_seq_len = med_seq_len = max_seq_len = 0
+
+        if name_order:
+            selected = self.take_seqs(name_order)
+        else:
+            name_order = list(self.names)
+            selected = self
+
+        # Stylise each character in each sequence
+        gaps = "".join(selected.moltype.gaps)
+        template = '<span class="%s">%%s</span>'
+        styled_seqs = defaultdict(list)
+        max_truncated_len = 0
+        for name in name_order:
+            sequence = str(self.named_seqs[name])[:limit]
+            seq_len = len(sequence)
+            max_truncated_len = max(seq_len, max_truncated_len)
+            start_gap = re.search(f"^[{gaps}]+", sequence)
+            end_gap = re.search(f"[{gaps}]+$", sequence)
+            start = 0 if start_gap is None else start_gap.end()
+            end = seq_len if end_gap is None else end_gap.start()
+
+            seq = []
+            for i, char in enumerate(sequence):
+                if i < start or i >= end:
+                    style = f"terminal_ambig_{self.moltype.label}"
+                else:
+                    style = styles[char]
+                s = template % style
+                s = s % char
+                seq.append(s)
+
+            styled_seqs[name] = seq
+
+        # Ensure all sublists are of same length
+        for name in styled_seqs:
+            if len(styled_seqs[name]) < max_truncated_len:
+                styled_seqs[name].extend(
+                    [""] * (max_truncated_len - len(styled_seqs[name]))
+                )
+
+        # Make html table
+        seqs = array([styled_seqs[n] for n in name_order], dtype="O")
+        table = ["<table>"]
+        seq_ = "<td>%s</td>"
+        label_ = '<td class="label">%s</td>'
+        num_row_ = '<tr class="num_row"><td></td><td><b>{:,d}</b></td></tr>'
+        for i in range(0, max_truncated_len, wrap):
+            table.append(num_row_.format(i))
+            seqblock = seqs[:, i : i + wrap].tolist()
+            for n, s in zip(name_order, seqblock):
+                s = "".join(s)
+                # Filter out rows that are empty (due to combination of shorter sequences + wrapping)
+                if len(s) > 0:
+                    row = "".join([label_ % n, seq_ % s])
+                    table.append(f"<tr>{row}</tr>")
+        table.append("</table>")
+        if (
+            limit
+            and limit < len(selected)
+            or name_order
+            and len(name_order) < len(selected.names)
+        ):
+            summary = (
+                "%s x {min=%s, median=%s, max=%s} (truncated to %s x %s) %s sequence collection"
+            ) % (
+                self.num_seqs,
+                min_seq_len,
+                med_seq_len,
+                max_seq_len,
+                len(name_order) if name_order else len(selected.names),
+                limit if limit else len(selected),
+                selected.moltype.label,
+            )
+        else:
+            summary = ("%s x {min=%s, median=%s, max=%s} %s sequence collection") % (
+                self.num_seqs,
+                min_seq_len,
+                med_seq_len,
+                max_seq_len,
+                selected.moltype.label,
+            )
+
+        text = [
+            "<style>",
+            ".c3align table {margin: 10px 0;}",
+            ".c3align td { border: none !important; text-align: left !important; }",
+            ".c3align tr:not(.num_row) td span {margin: 0 2px;}",
+            ".c3align tr:nth-child(even) {background: #f7f7f7;}",
+            ".c3align .num_row {background-color:rgba(161, 195, 209, 0.5) !important; border-top: solid 1px black; }",
+            ".c3align .label { font-size: %dpt ; text-align: right !important; "
+            "color: black !important; padding: 0 4px; display: table-cell !important; "
+            "font-weight: normal !important; }" % font_size,
+            "\n".join([".c3align " + style for style in css]),
+            "</style>",
+            '<div class="c3align">',
+            "\n".join(table),
+            f"<p><i>{summary}</i></p>",
+            "</div>",
+        ]
+        return "\n".join(text)
+
 
 @total_ordering
 class Aligned:
@@ -2210,7 +2426,7 @@ class Aligned:
 
         return self.__class__(new_map, new_seq)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.map!r} of {self.data!r}"
 
     def with_termini_unknown(self):
@@ -2886,7 +3102,7 @@ class AlignmentI(object):
 
         return names, output
 
-    def _repr_html_(self):
+    def _repr_html_(self) -> str:
         settings = self._repr_policy.copy()
         env_vals = get_setting_from_environ(
             "COGENT3_ALIGNMENT_REPR_POLICY",
@@ -2902,14 +3118,14 @@ class AlignmentI(object):
 
     def to_html(
         self,
-        name_order=None,
-        wrap=60,
-        limit=None,
-        ref_name="longest",
-        colors=None,
-        font_size=12,
-        font_family="Lucida Console",
-    ):
+        name_order: Optional[typing.Sequence[str]] = None,
+        wrap: int = 60,
+        limit: Optional[int] = None,
+        ref_name: str = "longest",
+        colors: Optional[Mapping[str, str]] = None,
+        font_size: int = 12,
+        font_family: str = "Lucida Console",
+    ) -> str:
         """returns html with embedded styles for sequence colouring
 
         Parameters
@@ -2933,16 +3149,28 @@ class AlignmentI(object):
         font_family
             string denoting font family
 
-        To display in jupyter notebook:
+        Examples
+        ---------
+
+        In a jupyter notebook, this code is used to provide the representation.
+
+        .. code-block:: python
+
+            aln # is rendered by jupyter
 
-            >>> from IPython.core.display import HTML
-            >>> HTML(aln.to_html())
+        You can directly use the result for display in a notebook as
+
+        .. code-block:: python
+
+            from IPython.core.display import HTML
+            HTML(aln.to_html())
         """
         css, styles = self.moltype.get_css_style(
             colors=colors, font_size=font_size, font_family=font_family
         )
         if name_order:
             selected = self.take_seqs(name_order)
+            name_order = list(name_order)
         else:
             name_order = list(self.names)
             ref_name = ref_name or "longest"
@@ -4049,7 +4277,7 @@ class ArrayAlignment(AlignmentI, _SequenceCollectionBase):
             result.append(">" + str(l) + "\n" + "".join(seq2str(s)))
         return "\n".join(result) + "\n"
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         seqs = []
         limit = 10
         delimiter = ""
@@ -4691,7 +4919,7 @@ class Alignment(AlignmentI, SequenceCollection):
 
         return new
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         seqs = []
         limit = 10
         delimiter = ""


Test patch:
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index dcb2ec9cd..8843e604f 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -859,10 +859,6 @@ class SequenceCollectionBaseTests(object):
 
     def test_set_wrap_affects_repr_html(self):
         """the wrap argument affects the number of columns"""
-        if self.Class == SequenceCollection:
-            # this class does not have this method
-            return
-
         # indirectly tested via counting number of occurrences of 'class="label"'
         seqs = self.Class({"a": "AAAAA", "b": "AAA--"})
         orig = seqs._repr_html_()
@@ -876,8 +872,8 @@ class SequenceCollectionBaseTests(object):
         os.environ[env_name] = "wrap=2"
         seqs = self.Class({"a": "AAAAA", "b": "AAA--"})
         got = seqs._repr_html_()
-        self.assertEqual(got.count(token), 3 * orig.count(token))
         os.environ.pop(env_name, None)
+        self.assertEqual(got.count(token), 3 * orig.count(token))
 
     def test_get_seq_entropy(self):
         """get_seq_entropy should get entropy of each seq"""
@@ -3637,3 +3633,45 @@ def test_quick_tree(cls, calc, brca1_data):
         if not edge.is_root()
     }
     assert types == {float}
+
+
+def test_sequence_collection_repr():
+    data = {
+        "ENSMUSG00000056468": "GCCAGGGGGGAAAGGGAGAA",
+        "ENSMUSG00000039616": "GCCCTTCAAATTT",
+    }
+    seqs = SequenceCollection(data=data, moltype=DNA)
+    assert (
+        repr(seqs)
+        == "2x (ENSMUSG00000039616[GCCCTTCAAAT...], ENSMUSG00000056468[GCCAGGGGGGA...]) dna seqcollection"
+    )
+
+    data = {
+        "ENSMUSG00000039616": "GCCCTTCAAATTT",
+        "ENSMUSG00000056468": "GCCAGGGGGGAAAGGGAGAA",
+    }
+    seqs = SequenceCollection(data=data, moltype=DNA)
+    assert (
+        repr(seqs)
+        == "2x (ENSMUSG00000039616[GCCCTTCAAAT...], ENSMUSG00000056468[GCCAGGGGGGA...]) dna seqcollection"
+    )
+
+    data = {
+        "a": "TCGAT",
+    }
+    seqs = SequenceCollection(data=data, moltype=DNA)
+    assert repr(seqs) == "1x (a[TCGAT]) dna seqcollection"
+
+    data = {
+        "a": "A" * 11,
+        "b": "B" * 3,
+        "c": "C" * 3,
+        "d": "D" * 11,
+        "e": "E" * 8,
+    }
+    seqs = SequenceCollection(data=data, moltype=ASCII)
+    assert repr(seqs) == "5x (b[BBB], ..., d[DDDDDDDDDDD...]) text seqcollection"
+
+    data = {}
+    seqs = SequenceCollection(data=data, moltype=BYTES)
+    assert repr(seqs) == "0x () bytes seqcollection"