diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index 15b1ff0f0e..22135f3042 100644 --- a/python/CHANGELOG.rst +++ b/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`. diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index ed9df8ea90..b05848c2d5 100644 --- a/python/tests/test_highlevel.py +++ b/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"{table.capitalize()}" 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) diff --git a/python/tests/test_util.py b/python/tests/test_util.py index b9568db429..a28256af1d 100644 --- a/python/tests/test_util.py +++ b/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║ +╚══╧══╧══╧══╝ +""" + ) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 2550b66429..51717b2a5b 100644 --- a/python/tskit/trees.py +++ b/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 diff --git a/python/tskit/util.py b/python/tskit/util.py index 77d10fc228..8d055440b2 100644 --- a/python/tskit/util.py +++ b/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"""