From 8ce88b64655f13a3cca3c6b991c6c4ad1024f0c8 Mon Sep 17 00:00:00 2001
From: LTLA <infinite.monkeys.with.keyboards@gmail.com>
Date: Fri, 27 Oct 2023 16:48:49 -0700
Subject: [PATCH] Added utility for pretty-printing of a truncated list.

This is equivalent to Bioconductor's coolcat function.
---
 src/biocutils/__init__.py             |  1 +
 src/biocutils/print_truncated_list.py | 51 +++++++++++++++++++++++++++
 tests/test_print_truncated_list.py    | 10 ++++++
 3 files changed, 62 insertions(+)
 create mode 100644 src/biocutils/print_truncated_list.py
 create mode 100644 tests/test_print_truncated_list.py

diff --git a/src/biocutils/__init__.py b/src/biocutils/__init__.py
index 92aaf8f..7e59c73 100644
--- a/src/biocutils/__init__.py
+++ b/src/biocutils/__init__.py
@@ -23,3 +23,4 @@
 from .subset import subset
 from .is_list_of_type import is_list_of_type
 from .normalize_subscript import normalize_subscript
+from .print_truncated_list import print_truncated_list
diff --git a/src/biocutils/print_truncated_list.py b/src/biocutils/print_truncated_list.py
new file mode 100644
index 0000000..494890c
--- /dev/null
+++ b/src/biocutils/print_truncated_list.py
@@ -0,0 +1,51 @@
+from typing import Optional, Sequence, Callable
+
+
+def print_truncated_list(x: Sequence, truncated_to = 3, full_threshold = 10, transform: Optional[Callable] = None, sep: str = ", ", include_brackets: bool = True) -> str:
+    """
+    Pretty-print a truncated list, replacing the middle elements with an
+    ellipsis if there are too many. This provides a useful preview of an
+    object without spewing out all of its contents on the screen.
+
+    Args:
+        x: List or some other sequence to be printed.
+
+        truncated_to:
+            Number of elements to truncate to, at the start and end of the
+            sequence. This should be less than half of ``full_threshold``.
+
+        full_threshold:
+            Threshold on the number of elements, below which the list is
+            shown in its entirety.
+
+        transform:
+            Optional transformation to apply to the elements of ``x``
+            after truncation but before printing.
+
+        sep:
+            Separator between elements in the printed list.
+
+        include_brackets:
+            Whether to include the start/end brackets.
+
+    Returns:
+        String containing the pretty-printed truncated list. 
+    """
+    collected = []
+    if transform is None:
+        transform = lambda y : y
+
+    if len(x) > full_threshold and len(x) > truncated_to * 2:
+        for i in range(truncated_to):
+            collected.append(repr(transform(x[i])))
+        collected.append("...")
+        for i in range(truncated_to, 0, -1):
+            collected.append(repr(transform(x[len(x) - i])))
+    else:
+        for c in x:
+            collected.append(repr(transform(c)))
+
+    output = sep.join(collected)
+    if include_brackets:
+        output = "[" + output + "]"
+    return output
diff --git a/tests/test_print_truncated_list.py b/tests/test_print_truncated_list.py
new file mode 100644
index 0000000..94f062b
--- /dev/null
+++ b/tests/test_print_truncated_list.py
@@ -0,0 +1,10 @@
+from biocutils import print_truncated_list
+
+
+def test_print_truncated_list():
+    assert print_truncated_list(range(6)) == repr(list(range(6)))
+    assert print_truncated_list(range(10)) == repr(list(range(10)))
+    assert print_truncated_list(range(200)) == "[0, 1, 2, ..., 197, 198, 199]"
+    assert print_truncated_list(["A", "B", "C", "D", "E", "F"], transform=lambda x : "foo_" + x) == "['foo_A', 'foo_B', 'foo_C', 'foo_D', 'foo_E', 'foo_F']"
+    assert print_truncated_list(["A", "B", "C", "D", "E", "F"], truncated_to=2, full_threshold=5, transform=lambda x : "foo_" + x) == "['foo_A', 'foo_B', ..., 'foo_E', 'foo_F']"
+    assert print_truncated_list(range(200), sep=" ", include_brackets=False) == "0 1 2 ... 197 198 199"