From 5789caa3d2f206027d339eef097884cddcb39a49 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:24:56 -0700 Subject: [PATCH] REF Deduplicate pretty output color formatting Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 026dff00-745e-42fb-9748-5555ee006073 --- pyrit/output/_formatting.py | 26 ++++++++++++++++++++++ pyrit/output/attack_result/pretty.py | 19 ++-------------- pyrit/output/conversation/pretty.py | 19 ++-------------- pyrit/output/scenario_result/pretty.py | 19 ++-------------- pyrit/output/score/pretty.py | 21 +++--------------- pyrit/output/scorer/pretty.py | 19 ++-------------- tests/unit/output/test_formatting.py | 30 ++++++++++++++++++++++++++ 7 files changed, 67 insertions(+), 86 deletions(-) create mode 100644 pyrit/output/_formatting.py create mode 100644 tests/unit/output/test_formatting.py diff --git a/pyrit/output/_formatting.py b/pyrit/output/_formatting.py new file mode 100644 index 0000000000..3413ea320e --- /dev/null +++ b/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" diff --git a/pyrit/output/attack_result/pretty.py b/pyrit/output/attack_result/pretty.py index b444472ed2..c31381fd88 100644 --- a/pyrit/output/attack_result/pretty.py +++ b/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, diff --git a/pyrit/output/conversation/pretty.py b/pyrit/output/conversation/pretty.py index ad5276b7f6..4cad444872 100644 --- a/pyrit/output/conversation/pretty.py +++ b/pyrit/output/conversation/pretty.py @@ -7,6 +7,7 @@ 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 @@ -14,7 +15,7 @@ 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. diff --git a/pyrit/output/scenario_result/pretty.py b/pyrit/output/scenario_result/pretty.py index 0584643ea5..44e677ca66 100644 --- a/pyrit/output/scenario_result/pretty.py +++ b/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. diff --git a/pyrit/output/score/pretty.py b/pyrit/output/score/pretty.py index 1f86e55fdd..97b046b467 100644 --- a/pyrit/output/score/pretty.py +++ b/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. diff --git a/pyrit/output/scorer/pretty.py b/pyrit/output/scorer/pretty.py index 2e0ddcbfe1..3b8e1594e9 100644 --- a/pyrit/output/scorer/pretty.py +++ b/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: diff --git a/tests/unit/output/test_formatting.py b/tests/unit/output/test_formatting.py new file mode 100644 index 0000000000..7e4c86f3c0 --- /dev/null +++ b/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