unit test for the truncate_output method

ersilia/cli/commands/run.py

@@ -1,56 +1,14 @@
-import json
 import types
 
 import click
 
 from ... import ErsiliaModel
 from ...core.session import Session
-from ...utils.terminal import print_result_table
+from ...utils.terminal import print_result_table, truncate_output
 from .. import echo
 from . import ersilia_cli
 
 
-def truncate_output(output, max_items=10, max_chars=500):
-    """
-    Truncates long outputs for better readability.
-
-    Parameters
-    ----------
-    output : Any
-        The output to process and truncate.
-    max_items : int, optional
-        Maximum number of items to display for arrays/lists or dictionary keys.
-    max_chars : int, optional
-        Maximum number of characters to display for strings.
-
-    Returns
-    -------
-    str
-        The truncated output as a formatted string.
-    """
-    if isinstance(output, list):
-        if len(output) > max_items:
-            return (
-                f"{output[:max_items]} ... (and {len(output) - max_items} more items)"
-            )
-        return str(output)
-    elif isinstance(output, dict):
-        formatted = json.dumps(output, indent=4)
-        lines = formatted.splitlines()
-        if len(lines) > max_items:
-            return (
-                "\n".join(lines[:max_items])
-                + f"\n... (and {len(lines) - max_items} more lines)"
-            )
-        return formatted
-    elif isinstance(output, str):
-        if len(output) > max_chars:
-            return f"{output[:max_chars]}... (truncated)"
-        return output
-    else:
-        return str(output)
-
-
 def run_cmd():
     """
     Runs a specified model.

ersilia/utils/terminal.py

@@ -159,6 +159,51 @@ def yes_no_input(prompt, default_answer, timeout=5):
         return True
 
 
+def truncate_output(output, max_items=10, max_chars=500):
+    """
+    Truncates long outputs for better readability.
+
+    Parameters
+    ----------
+    output : Any
+        The output to process and truncate.
+    max_items : int, optional
+        Maximum number of items to display for arrays/lists or dictionary keys.
+    max_chars : int, optional
+        Maximum number of characters to display for strings.
+
+    Returns
+    -------
+    str
+        The truncated output as a formatted string.
+    """
+    if isinstance(output, list):
+        if len(output) > max_items:
+            return (
+                f"{output[:max_items]} ... (and {len(output) - max_items} more items)"
+            )
+        return str(output)
+    if isinstance(output, dict):
+        if len(output) > max_items:
+            keys = list(output.keys())[:max_items]
+            truncated = {key: output[key] for key in keys}
+            remaining_lines = len(output) - max_items
+            return f"{json.dumps(truncated, indent=4)}\n... (and {remaining_lines} more lines)"
+        return output
+
+    elif isinstance(output, str):
+        if len(output) > max_chars:
+            suffix = " ... (truncated)"
+            truncated = output[
+                : max_chars - len(suffix) - 1
+            ].rstrip()  # Adjust for space
+            return f"{truncated} {suffix}"
+        return output
+
+    else:
+        return str(output)
+
+
 def print_result_table(data):
     """
     Print a result table from CSV or JSON-like data.

test/test_utils.py

@@ -0,0 +1,42 @@
+import pytest
+from ersilia.utils.terminal import truncate_output
+
+def test_truncate_output_list():
+    # Test truncation for a long list
+    output = list(range(20))  # List with 20 items
+    truncated = truncate_output(output, max_items=10)
+    assert truncated == "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] ... (and 10 more items)"
+
+def test_truncate_output_dict():
+    # Test truncation for a long dictionary
+    output = {f"key{i}": i for i in range(20)}  # Dictionary with 20 key-value pairs
+    truncated = truncate_output(output, max_items=5)
+    assert truncated.startswith("{\n    \"key0\": 0,")
+    assert truncated.endswith("... (and 15 more lines)")
+
+
+
+def test_truncate_output_short_list():
+    # Test for a short list that doesn't need truncation
+    output = [1, 2, 3]
+    truncated = truncate_output(output, max_items=10)
+    assert truncated == "[1, 2, 3]"
+
+def test_truncate_output_short_dict():
+    # Test for a short dictionary that doesn't need truncation
+    output = {"key1": 1, "key2": 2}
+    truncated = truncate_output(output, max_items=10)
+    assert "key1" in truncated
+    assert "key2" in truncated
+
+def test_truncate_output_short_string():
+    # Test for a short string that doesn't need truncation
+    output = "Short string"
+    truncated = truncate_output(output, max_chars=50)
+    assert truncated == "Short string"
+
+def test_truncate_output_other_types():
+    # Test for non-list, non-dict, non-string types
+    output = 12345
+    truncated = truncate_output(output)
+    assert truncated == "12345"