From e7eb96d39796b0cba54aafd510355ee597035255 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Fri, 23 Oct 2026 12:06:58 +0530 Subject: [PATCH] keep line breaks when wrapping image text --- pyrit/converter/base_image_text_converter.py | 10 +- .../converter/test_image_text_line_breaks.py | 92 +++++++++++++++++++ 2 files changed, 99 insertions(+), 3 deletions(-) create mode 100644 tests/unit/converter/test_image_text_line_breaks.py diff --git a/pyrit/converter/base_image_text_converter.py b/pyrit/converter/base_image_text_converter.py index 8aa60ba243..1e9a5806fc 100644 --- a/pyrit/converter/base_image_text_converter.py +++ b/pyrit/converter/base_image_text_converter.py @@ -151,7 +151,7 @@ def _fit_font_to_box( def _wrap_text(self, *, text: str, font: FreeTypeFont, max_width: int) -> list[str]: """ - Word-wrap text to fit within max_width pixels. + Word-wrap text to fit within max_width pixels, keeping the text's own line breaks. Args: text (str): The text to wrap. @@ -166,8 +166,12 @@ def _wrap_text(self, *, text: str, font: FreeTypeFont, max_width: int) -> list[s bbox = draw.textbbox((0, 0), string.ascii_letters, font=font) avg_char_width = (bbox[2] - bbox[0]) / len(string.ascii_letters) max_chars = max(1, int(max_width / avg_char_width)) - wrapped = textwrap.fill(text, width=max_chars) - return wrapped.split("\n") + # Wrap each line on its own; textwrap would otherwise turn line breaks into spaces. + return [ + wrapped_line + for line in text.splitlines() or [""] + for wrapped_line in textwrap.wrap(line, width=max_chars) or [""] + ] def _get_line_height(self, *, font: FreeTypeFont) -> int: """ diff --git a/tests/unit/converter/test_image_text_line_breaks.py b/tests/unit/converter/test_image_text_line_breaks.py new file mode 100644 index 0000000000..06c4252951 --- /dev/null +++ b/tests/unit/converter/test_image_text_line_breaks.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Text drawn onto images keeps its own line breaks.""" + +from pathlib import Path +from typing import Literal +from unittest.mock import patch + +import pytest +from PIL import Image + +from pyrit.converter import AddImageTextConverter, AddTextImageConverter, GridCompositeConverter +from pyrit.converter.base_image_text_converter import _BaseImageTextConverter + +ConverterKind = Literal["image_text_fixed", "image_text_auto", "text_image", "grid"] + + +@pytest.fixture +def wide_image(tmp_path: Path) -> str: + path = str(tmp_path / "wide.png") + Image.new("RGB", (1600, 800), color=(200, 200, 200)).save(path) + return path + + +@pytest.fixture +def innocuous_images(tmp_path: Path) -> list[str]: + paths = [] + for index in range(3): + path = str(tmp_path / f"innocuous_{index}.png") + Image.new("RGB", (120, 90), color=(90, 90, 90)).save(path) + paths.append(path) + return paths + + +async def _drawn_lines_async( + *, kind: ConverterKind, text: str, wide_image: str, innocuous_images: list[str] +) -> list[str]: + with patch.object( + _BaseImageTextConverter, + "_draw_text_overlay", + autospec=True, + side_effect=_BaseImageTextConverter._draw_text_overlay, + ) as draw: + if kind == "image_text_fixed": + await AddImageTextConverter(img_to_add=Path(wide_image)).convert_async(prompt=text, input_type="text") + elif kind == "image_text_auto": + converter = AddImageTextConverter(img_to_add=Path(wide_image), font_size=(10, 20)) + await converter.convert_async(prompt=text, input_type="text") + elif kind == "text_image": + await AddTextImageConverter(text_to_add=text).convert_async(prompt=wide_image, input_type="image_path") + else: + converter = GridCompositeConverter(innocuous_images=innocuous_images, tile_size=(1200, 600)) + await converter.convert_async(prompt=text, input_type="text") + lines = draw.call_args.kwargs["lines"] + assert isinstance(lines, list) + return [str(line) for line in lines] + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("kind", ["image_text_fixed", "image_text_auto", "text_image", "grid"]) +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("Step 1\nStep 2", ["Step 1", "Step 2"]), + ("Step 1\r\nStep 2", ["Step 1", "Step 2"]), + ("Intro\n\nDone", ["Intro", "", "Done"]), + ("Just one line", ["Just one line"]), + ], +) +async def test_line_breaks_are_drawn_as_lines_async( + *, kind: ConverterKind, text: str, expected: list[str], wide_image: str, innocuous_images: list[str] +) -> None: + lines = await _drawn_lines_async(kind=kind, text=text, wide_image=wide_image, innocuous_images=innocuous_images) + + assert lines == expected + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("kind", ["image_text_fixed", "text_image"]) +async def test_long_lines_still_wrap_after_a_break_async( + *, kind: ConverterKind, wide_image: str, innocuous_images: list[str] +) -> None: + long_line = " ".join(["word"] * 400) + + lines = await _drawn_lines_async( + kind=kind, text=f"Title\n{long_line}", wide_image=wide_image, innocuous_images=innocuous_images + ) + + assert lines[0] == "Title" + assert len(lines) > 2 + assert " ".join(lines[1:]) == long_line