From 6dec77670097e455193dd1c186fd65ccd4721731 Mon Sep 17 00:00:00 2001 From: Nakshatra Mote Date: Mon, 15 Jun 2026 14:44:16 +0530 Subject: [PATCH] Fix matplotlib warnings and TclError in image generator --- agentic_security/probe_data/image_generator.py | 10 +++++++++- agentic_security/probe_data/test_image_generator.py | 9 +++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/agentic_security/probe_data/image_generator.py b/agentic_security/probe_data/image_generator.py index d986651..c417c81 100644 --- a/agentic_security/probe_data/image_generator.py +++ b/agentic_security/probe_data/image_generator.py @@ -1,8 +1,12 @@ import base64 import io +import re import httpx -import matplotlib.pyplot as plt +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 from cache_to_disk import cache_to_disk from tqdm import tqdm @@ -49,6 +53,10 @@ def generate_image(prompt: str, variant: int = 0) -> bytes: Returns: bytes: The image data in JPG format. """ + # Sanitize prompt: replace non-renderable whitespace characters (tabs, etc.) + # with spaces to avoid matplotlib UserWarning about missing glyphs. + prompt = re.sub(r"[\t\r\x0b\x0c]", " ", prompt) + # Create a matplotlib figure fig, ax = plt.subplots(figsize=(6, 4)) diff --git a/agentic_security/probe_data/test_image_generator.py b/agentic_security/probe_data/test_image_generator.py index 30304c7..ce852d2 100644 --- a/agentic_security/probe_data/test_image_generator.py +++ b/agentic_security/probe_data/test_image_generator.py @@ -39,3 +39,12 @@ def test_generate_image_dataset(mock_generate_image): assert isinstance(image_datasets[0], ImageProbeDataset) assert image_datasets[0].test_dataset.dataset_name == test_dataset_name assert image_datasets[0].image_prompts[0] == b"dummy_image_bytes" + + +def test_generate_image_with_special_whitespace(): + """Test that prompts with tab and other non-renderable whitespace don't raise warnings.""" + prompt_with_tabs = "Hello\tWorld\tTest" + image_bytes = generate_image(prompt_with_tabs, 0) + + assert isinstance(image_bytes, bytes) + assert len(image_bytes) > 0