diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 511c68de..42a436c9 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -2,6 +2,7 @@ from __future__ import annotations import textwrap from enum import Enum +from typing import Any import numpy as np @@ -151,30 +152,36 @@ def resolve_color( return get_color_by_index(color=color, idx=idx) -def wrap_text(text: str, max_line_length=None) -> list[str]: +def wrap_text(text: Any, max_line_length=None) -> list[str]: """ - Wraps text to the specified maximum line length, respecting existing newlines. - Uses the textwrap library for robust text wrapping. + Wrap `text` to the specified maximum line length, respecting existing + newlines. Falls back to str() if `text` is not already a string. Args: - text (str): The text to wrap. + text (Any): The text (or object) to wrap. + max_line_length (int | None): Maximum width for each wrapped line. Returns: - List[str]: A list of text lines after wrapping. + list[str]: Wrapped lines. """ if not text: return [""] + if not isinstance(text, str): + text = str(text) + if max_line_length is None: return text.splitlines() or [""] + if max_line_length <= 0: + raise ValueError("max_line_length must be a positive integer") + paragraphs = text.split("\n") - all_lines = [] + all_lines: list[str] = [] for paragraph in paragraphs: - if not paragraph: - # Keep empty lines + if paragraph == "": all_lines.append("") continue @@ -186,12 +193,9 @@ def wrap_text(text: str, max_line_length=None) -> list[str]: drop_whitespace=True, ) - if wrapped: - all_lines.extend(wrapped) - else: - all_lines.append("") + all_lines.extend(wrapped or [""]) - return all_lines if all_lines else [""] + return all_lines or [""] def validate_labels(labels: list[str] | None, detections: Detections): diff --git a/test/annotators/test_utils.py b/test/annotators/test_utils.py index d6abc381..3ab0f9b9 100644 --- a/test/annotators/test_utils.py +++ b/test/annotators/test_utils.py @@ -5,7 +5,7 @@ from contextlib import ExitStack as DoesNotRaise import numpy as np import pytest -from supervision.annotators.utils import ColorLookup, resolve_color_idx +from supervision.annotators.utils import ColorLookup, resolve_color_idx, wrap_text from supervision.detection.core import Detections from test.test_utils import mock_detections @@ -108,3 +108,67 @@ def test_resolve_color_idx( color_lookup=color_lookup, ) assert result == expected_result + + +@pytest.mark.parametrize( + "text, max_line_length, expected_result, exception", + [ + (None, None, [""], DoesNotRaise()), # text is None + ("", None, [""], DoesNotRaise()), # empty string + (" \t ", 3, [""], DoesNotRaise()), # whitespace-only (spaces + tab) + (12345, None, ["12345"], DoesNotRaise()), # plain integer + (-6789, None, ["-6789"], DoesNotRaise()), # negative integer + (np.int64(1000), None, ["1000"], DoesNotRaise()), # NumPy int64 + ([1, 2, 3], None, ["[1, 2, 3]"], DoesNotRaise()), # list to string + ( + "When you play the game of thrones, you win or you die.\nFear cuts deeper than swords.\nA mind needs books as a sword needs a whetstone.", # noqa: E501 + None, + [ + "When you play the game of thrones, you win or you die.", + "Fear cuts deeper than swords.", + "A mind needs books as a sword needs a whetstone.", + ], + DoesNotRaise(), + ), # Game-of-Thrones quotes, multiline + ("\n", None, [""], DoesNotRaise()), # single newline + ( + "valarmorghulisvalardoharis", + 6, + ["valarm", "orghul", "isvala", "rdohar", "is"], + DoesNotRaise(), + ), # long Valyrian phrase, wrapped + ( + "Winter is coming\nFire and blood", + 10, + [ + "Winter is", + "coming", + "Fire and", + "blood", + ], + DoesNotRaise(), + ), # mix of short/long with newline + ( + "What is dead may never die", + 0, + None, + pytest.raises(ValueError), + ), # width 0 - invalid + ( + "A Lannister always pays his debts", + -1, + None, + pytest.raises(ValueError), + ), # width -1 - invalid + (None, 10, [""], DoesNotRaise()), # text None, width set + ], +) +def test_wrap_text( + text: object, + max_line_length: int | None, + expected_result: list[str], + exception: Exception, +) -> None: + with exception: + result = wrap_text(text=text, max_line_length=max_line_length) + assert result == expected_result