MAINT Deduplicate pretty output color formatting (#2318)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 026dff00-745e-42fb-9748-5555ee006073

Author: romanlutz

pyrit/output/_formatting.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from colorama import Style
+
+
+class _PrettyPrinterMixin:
+    """Shared ANSI line formatting for pretty printers."""
+
+    _enable_colors: bool
+
+    def _format_colored(self, text: str, *colors: str) -> str:
+        """
+        Format text with color codes if colors are enabled.
+
+        Args:
+            text (str): The text to format.
+            *colors: Variable number of colorama color constants to apply.
+
+        Returns:
+            str: The formatted line with trailing newline.
+        """
+        if self._enable_colors and colors:
+            color_prefix = "".join(colors)
+            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
+        return f"{text}\n"

pyrit/output/attack_result/pretty.py

@@ -7,13 +7,14 @@
 from colorama import Back, Fore, Style
 
 from pyrit.models import AttackOutcome, AttackResult, ConversationType, Message, Score
+from pyrit.output._formatting import _PrettyPrinterMixin
 from pyrit.output.attack_result.base import AttackResultPrinterBase
 from pyrit.output.conversation.pretty import PrettyConversationPrinter
 from pyrit.output.score.pretty import PrettyScorePrinter
 from pyrit.output.sink import Sink
 
 
-class PrettyAttackResultPrinter(AttackResultPrinterBase):
+class PrettyAttackResultPrinter(_PrettyPrinterMixin, AttackResultPrinterBase):
     """
     Pretty printer for attack results with ANSI-colored formatting.
 
@@ -68,22 +69,6 @@ def __init__(
             blur_radius=blur_radius,
         )
 
-    def _format_colored(self, text: str, *colors: str) -> str:
-        """
-        Format text with color codes if colors are enabled.
-
-        Args:
-            text (str): The text to format.
-            *colors: Variable number of colorama color constants to apply.
-
-        Returns:
-            str: The formatted line with trailing newline.
-        """
-        if self._enable_colors and colors:
-            color_prefix = "".join(colors)
-            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
-        return f"{text}\n"
-
     async def render_async(
         self,
         result: AttackResult,

pyrit/output/conversation/pretty.py

@@ -7,14 +7,15 @@
 from colorama import Fore, Style
 
 from pyrit.models import Message, MessagePiece, Score
+from pyrit.output._formatting import _PrettyPrinterMixin
 from pyrit.output.conversation.base import ConversationPrinterBase
 from pyrit.output.score.pretty import PrettyScorePrinter
 from pyrit.output.sink import Sink
 
 logger = logging.getLogger(__name__)
 
 
-class PrettyConversationPrinter(ConversationPrinterBase):
+class PrettyConversationPrinter(_PrettyPrinterMixin, ConversationPrinterBase):
     """
     Pretty printer for conversation message histories with ANSI-colored formatting.
 
@@ -174,22 +175,6 @@ async def render_async(
 
         return "".join(lines)
 
-    def _format_colored(self, text: str, *colors: str) -> str:
-        """
-        Format text with color codes if colors are enabled.
-
-        Args:
-            text (str): The text to format.
-            *colors: Variable number of colorama color constants to apply.
-
-        Returns:
-            str: The formatted line with trailing newline.
-        """
-        if self._enable_colors and colors:
-            color_prefix = "".join(colors)
-            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
-        return f"{text}\n"
-
     def _render_wrapped_text(self, text: str, color: str) -> str:
         """
         Render text with proper wrapping and indentation, preserving newlines.

pyrit/output/scenario_result/pretty.py

@@ -6,12 +6,13 @@
 from colorama import Fore, Style
 
 from pyrit.models import AttackOutcome, ScenarioResult
+from pyrit.output._formatting import _PrettyPrinterMixin
 from pyrit.output.scenario_result.base import ScenarioResultPrinterBase
 from pyrit.output.scorer.base import ScorerPrinterBase
 from pyrit.output.sink import Sink
 
 
-class PrettyScenarioResultPrinter(ScenarioResultPrinterBase):
+class PrettyScenarioResultPrinter(_PrettyPrinterMixin, ScenarioResultPrinterBase):
     """
     Pretty printer for scenario results with ANSI-colored formatting.
 
@@ -51,22 +52,6 @@ def __init__(
         self._scorer_printer = scorer_printer
         self._sort_groups_by_success_rate = sort_groups_by_success_rate
 
-    def _format_colored(self, text: str, *colors: str) -> str:
-        """
-        Format text with color codes if colors are enabled.
-
-        Args:
-            text (str): The text to format.
-            *colors: Variable number of colorama color constants to apply.
-
-        Returns:
-            str: The formatted line with trailing newline.
-        """
-        if self._enable_colors and colors:
-            color_prefix = "".join(colors)
-            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
-        return f"{text}\n"
-
     def _render_section_header(self, title: str) -> str:
         """
         Render a section header with visual separation.

pyrit/output/score/pretty.py

@@ -3,14 +3,15 @@
 
 import textwrap
 
-from colorama import Fore, Style
+from colorama import Fore
 
 from pyrit.models import Score
+from pyrit.output._formatting import _PrettyPrinterMixin
 from pyrit.output.base import PrinterBase
 from pyrit.output.sink import Sink
 
 
-class PrettyScorePrinter(PrinterBase):
+class PrettyScorePrinter(_PrettyPrinterMixin, PrinterBase):
     """
     Pretty printer for individual Score objects with ANSI-colored formatting.
 
@@ -36,22 +37,6 @@ def __init__(
         self._indent = " " * indent_size
         self._enable_colors = enable_colors
 
-    def _format_colored(self, text: str, *colors: str) -> str:
-        """
-        Format text with color codes if colors are enabled.
-
-        Args:
-            text (str): The text to format.
-            *colors: Variable number of colorama color constants to apply.
-
-        Returns:
-            str: The formatted line with trailing newline.
-        """
-        if self._enable_colors and colors:
-            color_prefix = "".join(colors)
-            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
-        return f"{text}\n"
-
     def _render_score(self, score: Score, indent_level: int = 3) -> str:
         """
         Render a single score with proper formatting.

pyrit/output/scorer/pretty.py

@@ -6,11 +6,12 @@
 from colorama import Fore, Style
 
 from pyrit.models import ComponentIdentifier
+from pyrit.output._formatting import _PrettyPrinterMixin
 from pyrit.output.scorer.base import ScorerPrinterBase
 from pyrit.output.sink import Sink
 
 
-class PrettyScorerPrinter(ScorerPrinterBase):
+class PrettyScorerPrinter(_PrettyPrinterMixin, ScorerPrinterBase):
     """
     Pretty printer for scorer information with ANSI-colored formatting.
 
@@ -39,22 +40,6 @@ def __init__(self, *, sink: Sink | None = None, indent_size: int = 2, enable_col
         self._indent = " " * indent_size
         self._enable_colors = enable_colors
 
-    def _format_colored(self, text: str, *colors: str) -> str:
-        """
-        Format text with color codes if colors are enabled.
-
-        Args:
-            text (str): The text to format.
-            *colors: Variable number of colorama color constants to apply.
-
-        Returns:
-            str: The formatted line with trailing newline.
-        """
-        if self._enable_colors and colors:
-            color_prefix = "".join(colors)
-            return f"{color_prefix}{text}{Style.RESET_ALL}\n"
-        return f"{text}\n"
-
     def _get_quality_color(
         self, value: float, *, higher_is_better: bool, good_threshold: float, bad_threshold: float
     ) -> str:

tests/unit/output/test_formatting.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import pytest
+from colorama import Fore, Style
+
+from pyrit.output._formatting import _PrettyPrinterMixin
+
+
+class _TestPrettyPrinter(_PrettyPrinterMixin):
+    def __init__(self, *, enable_colors: bool) -> None:
+        self._enable_colors = enable_colors
+
+
+@pytest.mark.parametrize(
+    "enable_colors,colors,expected",
+    [
+        (True, (Style.BRIGHT, Fore.RED), f"{Style.BRIGHT}{Fore.RED}text{Style.RESET_ALL}\n"),
+        (False, (Style.BRIGHT, Fore.RED), "text\n"),
+        (True, (), "text\n"),
+    ],
+)
+def test_format_colored_preserves_line_output(
+    enable_colors: bool,
+    colors: tuple[str, ...],
+    expected: str,
+) -> None:
+    printer = _TestPrettyPrinter(enable_colors=enable_colors)
+
+    assert printer._format_colored("text", *colors) == expected