Minor refactoring stage:list command (#5498)

Moving text truncation to utils and adding tests for it

dvc/command/stage.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Dict, Iterable, List
 from dvc.command import completion
 from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
 from dvc.utils.cli_parse import parse_params
+from dvc.utils.humanize import truncate_text
 
 if TYPE_CHECKING:
     from dvc.output.base import BaseOutput
@@ -14,7 +15,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 MAX_TEXT_LENGTH = 80
-ELLIPSIS = "…"
 
 
 def generate_description(stage: "Stage") -> str:
@@ -47,9 +47,7 @@ def prepare_description(
     stage: "Stage", max_length: int = MAX_TEXT_LENGTH
 ) -> str:
     desc = stage.short_description() or generate_description(stage)
-    if len(desc) > max_length:
-        return desc[: max_length - 1] + ELLIPSIS
-    return desc
+    return truncate_text(desc, max_length)
 
 
 def prepare_stages_data(

dvc/utils/humanize.py

@@ -25,3 +25,15 @@ def get_summary(stats):
         "{} file{} {}".format(num, "s" if num > 1 else "", state)
         for state, num in status
     )
+
+
+ELLIPSIS = "…"
+
+
+def truncate_text(
+    text: str, max_length: int, with_ellipsis: bool = True
+) -> str:
+    if with_ellipsis and len(text) > max_length:
+        return text[: max_length - 1] + ELLIPSIS
+
+    return text[:max_length]

tests/unit/utils/test_humanize.py

@@ -1,6 +1,8 @@
 from collections import OrderedDict
 
-from dvc.utils.humanize import get_summary
+import pytest
+
+from dvc.utils.humanize import get_summary, truncate_text
 
 
 def test_get_summary():
@@ -38,3 +40,35 @@ def test_get_summary():
     assert get_summary([]) == ""
     assert get_summary([("x", 0), ("y", [])]) == ""
     assert get_summary([("x", 1), ("y", [])]) == "1 file x"
+
+
+def test_truncate_text():
+    text = "lorem ipsum"
+    length = 5
+
+    truncated = truncate_text(text, length)
+    # length should not cross the max length
+    assert len(truncated) == length
+    assert truncated[:-1] == text[: length - 1]
+    # last character should be ellipsis
+    assert truncated[-1] == "…"
+
+    truncated = truncate_text(text, length, with_ellipsis=False)
+    # length should not cross the max length
+    assert len(truncated) == length
+    assert truncated == text[:length]
+
+
+@pytest.mark.parametrize("with_ellipsis", [True, False])
+def test_truncate_text_smaller_than_max_length(with_ellipsis):
+    text = "lorem ipsum"
+
+    # exact match as length
+    truncated = truncate_text(text, len(text), with_ellipsis=with_ellipsis)
+    assert len(truncated) == len(text)
+    assert truncated == text
+
+    # max_length > len(text)
+    truncated = truncate_text(text, len(text) + 1, with_ellipsis=with_ellipsis)
+    assert len(truncated) == len(text)
+    assert truncated == text