Add __repr__ to TreeSequence

python/CHANGELOG.rst

@@ -31,6 +31,9 @@
 - Added ``TreeSequence._repr_html_`` for use in jupyter notebooks.
   (:user:`benjeffery`, :issue:`872`, :pr:`923`)
 
+- Added ``TreeSequence.__repr__`` to display a summary for terminal usage.
+  (:user:`benjeffery`, :issue:`938`, :pr:`985`)
+
 **Breaking changes**
 
 - The argument to ``ts.dump`` and ``tskit.load`` has been renamed `file` from `path`.

python/tests/test_highlevel.py

@@ -34,6 +34,7 @@
 import pickle
 import platform
 import random
+import re
 import shutil
 import tempfile
 import unittest
@@ -1460,6 +1461,14 @@ def test_html_repr(self):
             for table in ts.tables.name_map:
                 assert f"<td>{table.capitalize()}</td>" in html
 
+    def test_repr(self):
+        for ts in get_example_tree_sequences():
+            s = repr(ts)
+            assert len(s) > 999
+            assert re.search(rf"║Trees *│ *{ts.num_trees}║", s)
+            for table in ts.tables.name_map:
+                assert re.search(rf"║{table.capitalize()} *│", s)
+
 
 class TestTreeSequenceMethodSignatures:
     ts = msprime.simulate(10, random_seed=1234)

python/tests/test_util.py

@@ -400,3 +400,37 @@ def test_obj_to_collapsed_html(obj, expected):
         util.obj_to_collapsed_html(obj, "Test", 1).replace(" ", "").replace("\n", "")
         == expected
     )
+
+
+def test_unicode_table():
+    assert (
+        util.unicode_table(
+            [["5", "6", "7", "8"], ["90", "10", "11", "12"]],
+            header=["1", "2", "3", "4"],
+        )
+        == """╔══╤══╤══╤══╗
+║1 │2 │3 │4 ║
+╠══╪══╪══╪══╣
+║5 │ 6│ 7│ 8║
+╟──┼──┼──┼──╢
+║90│10│11│12║
+╚══╧══╧══╧══╝
+"""
+    )
+
+    assert (
+        util.unicode_table(
+            [["1", "2", "3", "4"], ["5", "6", "7", "8"], ["90", "10", "11", "12"]],
+            title="TITLE",
+        )
+        == """╔═══════════╗
+║TITLE      ║
+╠══╤══╤══╤══╣
+║1 │ 2│ 3│ 4║
+╟──┼──┼──┼──╢
+║5 │ 6│ 7│ 8║
+╟──┼──┼──┼──╢
+║90│10│11│12║
+╚══╧══╧══╧══╝
+"""
+    )

python/tskit/trees.py

@@ -3334,6 +3334,33 @@ def dump_text(
                 )
                 print(row, file=provenances)
 
+    def __repr__(self):
+        ts_rows = [
+            ["Trees", str(self.num_trees)],
+            ["Sequence Length", str(self.sequence_length)],
+            ["Sample Nodes", str(self.num_samples)],
+            ["Total Size TODO", util.naturalsize(99999)],
+        ]
+        header = ["Table", "Rows", "Size", "Has Metadata"]
+        table_rows = []
+        for name, table in self.tables.name_map.items():
+            table_rows.append(
+                [
+                    str(s)
+                    for s in [
+                        name.capitalize(),
+                        table.num_rows,
+                        "TODO",
+                        "Yes"
+                        if hasattr(table, "metadata") and len(table.metadata) > 0
+                        else "No",
+                    ]
+                ]
+            )
+        return util.unicode_table(ts_rows, title="TreeSequence") + util.unicode_table(
+            table_rows, header=header
+        )
+
     def _repr_html_(self):
         """
         Called by jupyter notebooks to render a TreeSequence

python/tskit/util.py

@@ -287,6 +287,15 @@ def naturalsize(value):
 
 
 def obj_to_collapsed_html(d, name=None, open_depth=0):
+    """
+    Recursively make an HTML representation of python objects.
+
+    :param str name: Name for this object
+    :param int open_depth: By default sub-sections are collapsed. If this number is
+    non-zero the first layers up to open_depth will be opened.
+    :return: The HTML as a string
+    :rtype: str
+    """
     opened = "open" if open_depth > 0 else ""
     open_depth -= 1
     name = str(name) + ":" if name is not None else ""
@@ -316,6 +325,51 @@ def obj_to_collapsed_html(d, name=None, open_depth=0):
         return f"{name} {d}"
 
 
+def unicode_table(rows, title=None, header=None):
+    """
+    Convert a table (list of lists) of strings to a unicode table.
+
+    :param list[list[str]] rows: List of rows, each of which is a list of strings for
+    each cell. The first column will be left justified, the others right. Each row must
+    have the same number of cells.
+    :param str title: If specified the first output row will be a single cell
+    containing this string, left-justified. [optional]
+    :param list[str] header: Specifies a row above the main rows which will be in double
+    lined borders and left justified. Must be same length as each row. [optional]
+    :return: The table as a string
+    :rtype: str
+    """
+    if header is not None:
+        all_rows = [header] + rows
+    else:
+        all_rows = rows
+    widths = [
+        max(len(row[i_col]) for row in all_rows) for i_col in range(len(all_rows[0]))
+    ]
+    out = []
+    if title is not None:
+        w = sum(widths) + len(rows[1]) - 1
+        out += [
+            f"╔{'═' * w}╗\n" f"║{title.ljust(w)}║\n",
+            f"╠{'╤'.join('═' * w for w in widths)}╣\n",
+        ]
+    if header is not None:
+        out += [
+            f"╔{'╤'.join('═' * w for w in widths)}╗\n"
+            f"║{'│'.join(cell.ljust(w) for cell,w in zip(header,widths))}║\n",
+            f"╠{'╪'.join('═' * w for w in widths)}╣\n",
+        ]
+    out += [
+        f"╟{'┼'.join('─' * w for w in widths)}╢\n".join(
+            f"║{row[0].ljust(widths[0])}│"
+            + f"{'│'.join(cell.rjust(w) for cell, w in zip(row[1:], widths[1:]))}║\n"
+            for row in rows
+        ),
+        f"╚{'╧'.join('═' * w for w in widths)}╝\n",
+    ]
+    return "".join(out)
+
+
 def tree_sequence_html(ts):
     table_rows = "".join(
         f"""