diff --git a/ersilia/cli/commands/run.py b/ersilia/cli/commands/run.py index d0a0ae0b..c26e711d 100644 --- a/ersilia/cli/commands/run.py +++ b/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. diff --git a/ersilia/utils/terminal.py b/ersilia/utils/terminal.py index 2de80e29..67622f91 100644 --- a/ersilia/utils/terminal.py +++ b/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. diff --git a/test/test_utils.py b/test/test_utils.py new file mode 100644 index 00000000..462358f5 --- /dev/null +++ b/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"