diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 7eb642aec6..53136e15c8 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -345,6 +345,10 @@ def print_scenario_run_progress(*, run: ScenarioRunSummary, total_techniques: in bar_width = 30 filled = int(bar_width * techniques_done / effective_total) bar = "█" * filled + "░" * (bar_width - filled) + try: + bar.encode(sys.stdout.encoding or "utf-8") + except (LookupError, UnicodeEncodeError): + bar = "#" * filled + "-" * (bar_width - filled) parts.append(f"[{bar}] techniques: {techniques_done}/{effective_total} ({pct}%)") else: parts.append(f"techniques: {techniques_done}") diff --git a/pyrit/output/sink.py b/pyrit/output/sink.py index f7d9642475..e4549fbc30 100644 --- a/pyrit/output/sink.py +++ b/pyrit/output/sink.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import asyncio +import sys from abc import ABC, abstractmethod from pathlib import Path from typing import Literal @@ -41,7 +42,12 @@ async def write_async(self, data: str) -> None: Args: data (str): The text to print. """ - print(data, end="") + encoding = sys.stdout.encoding or "utf-8" + try: + data.encode(encoding) + except (LookupError, UnicodeEncodeError): + data = data.encode(encoding, errors="replace").decode(encoding) + sys.stdout.write(data) class FileSink(Sink): diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index c3fd62a7c2..84c18fa05b 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -389,6 +389,24 @@ def test_print_scenario_run_progress_with_known_totals(capsys): assert "attacks" not in captured.out +def test_print_scenario_run_progress_uses_ascii_for_limited_console(): + run = _make_run( + status=ScenarioRunState.IN_PROGRESS, + total_attacks=10, + completed_attacks=5, + objective_achieved_rate=30, + techniques_used=["s1"], + ) + stdout = MagicMock() + stdout.encoding = "cp1252" + + with patch.object(_output.sys, "stdout", stdout): + _output.print_scenario_run_progress(run=run, total_techniques=2) + + line = stdout.write.call_args.args[0] + assert "[###############---------------]" in line + + def test_print_scenario_run_progress_no_techniques(capsys): run = _make_run( status=ScenarioRunState.CREATED, diff --git a/tests/unit/output/test_sink.py b/tests/unit/output/test_sink.py index f4a5b85238..5859f401a3 100644 --- a/tests/unit/output/test_sink.py +++ b/tests/unit/output/test_sink.py @@ -31,6 +31,17 @@ async def test_stdout_sink_no_trailing_newline(capsys): assert captured.out == "line1line2" +async def test_stdout_sink_replaces_unsupported_unicode(): + stdout = MagicMock() + stdout.encoding = "cp1252" + sink = StdoutSink() + + with patch("pyrit.output.sink.sys.stdout", stdout): + await sink.write_async("📊─") + + stdout.write.assert_called_once_with("??") + + async def test_file_sink_writes_to_file(): with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as f: path = Path(f.name)