MAINT: make puzzled dispatch exhaustive (P311-03)

pyrit/converter/puzzled/puzzled_converter.py

@@ -5,6 +5,7 @@
 import json
 import pathlib
 import uuid
+from typing import TYPE_CHECKING, assert_never
 
 from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH
 from pyrit.converter.converter import Converter, ConverterResult
@@ -21,6 +22,9 @@
 from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, SeedPrompt
 from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS, PromptTarget
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 # Instruction for the optional clue-generation model. It produces the paper's "indirect
 # semantic description" (Ahn & Lee, arXiv:2508.01306): a short euphemistic definition that lets
 # a solver recover the word from its meaning without the word appearing verbatim.
@@ -192,13 +196,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
         if puzzle_type is PuzzleType.CROSSWORD and not crossword_symbol_map(words):
             puzzle_type = PuzzleType.ANAGRAM
 
-        rng = self._get_random_generator(stream="puzzle")
-        if puzzle_type is PuzzleType.WORD_SEARCH:
-            puzzle_body = build_word_search(words, rng)
-        elif puzzle_type is PuzzleType.ANAGRAM:
-            puzzle_body = build_anagram(words, rng)
-        else:
-            puzzle_body = build_crossword(words)
+        puzzle_body = self._build_puzzle_body(words=words, puzzle_type=puzzle_type)
 
         semantics = await self._semantic_clues_async(words)
         clues = "\n".join(self._clue_line(masked, semantics) for masked in mask_result.masked_words)
@@ -213,6 +211,29 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
 
         return ConverterResult(output_text=formatted_prompt, output_type="text")
 
+    def _build_puzzle_body(self, *, words: list[str], puzzle_type: PuzzleType) -> str:
+        """
+        Build the configured puzzle encoding and reject unhandled enum extensions.
+
+        Args:
+            words (list[str]): The masked words to encode.
+            puzzle_type (PuzzleType): The selected encoding after any safe fallback.
+
+        Returns:
+            str: The rendered puzzle body.
+        """
+        builder: Callable[..., str]
+        if puzzle_type is PuzzleType.WORD_SEARCH:
+            builder = build_word_search
+        elif puzzle_type is PuzzleType.ANAGRAM:
+            builder = build_anagram
+        elif puzzle_type is PuzzleType.CROSSWORD:
+            return build_crossword(words)
+        else:
+            assert_never(puzzle_type)
+
+        return builder(words, self._get_random_generator(stream="puzzle"))
+
     @staticmethod
     def _clue_line(masked: MaskedWord, semantics: dict[str, str]) -> str:
         """

tests/unit/converter/test_puzzled_converter.py

@@ -1,8 +1,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import inspect
+import subprocess
+import sys
 import unicodedata
 from collections import Counter
+from pathlib import Path
+from textwrap import dedent
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -166,6 +171,68 @@ async def test_crossword_falls_back_to_anagram_when_it_cannot_hide_the_word():
     assert "malware" not in result.output_text.lower()
 
 
+@pytest.mark.parametrize("extend_enum", [False, True], ids=["current-enum", "extended-enum"])
+def test_puzzle_dispatch_type_exhaustiveness(tmp_path: Path, extend_enum: bool) -> None:
+    enum_source = inspect.getsource(PuzzleType)
+    if extend_enum:
+        enum_source += '\n    FUTURE_KIND = "future_kind"\n'
+    probe = tmp_path / "dispatch.py"
+    converter_source = inspect.getsource(PuzzledConverter._build_puzzle_body)
+    probe.write_text(
+        "\n".join(
+            [
+                "import random",
+                "from collections.abc import Callable",
+                "from enum import Enum",
+                "from typing import assert_never",
+                enum_source,
+                dedent(
+                    """
+                    def build_word_search(words: list[str], rng: random.Random) -> str:
+                        return ""
+
+                    def build_anagram(words: list[str], rng: random.Random) -> str:
+                        return ""
+
+                    def build_crossword(words: list[str]) -> str:
+                        return ""
+
+                    class PuzzledConverter:
+                        def _get_random_generator(self, *, stream: str) -> random.Random:
+                            return random.Random()
+                    """
+                ),
+                converter_source,
+            ]
+        ),
+        encoding="utf-8",
+    )
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "ty",
+            "check",
+            "--project",
+            str(Path(__file__).resolve().parents[3]),
+            "--python-version",
+            "3.11",
+            str(probe),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=30,
+        check=False,
+    )
+    output = result.stdout + result.stderr
+    assert result.returncode == (1 if extend_enum else 0), output
+    if extend_enum:
+        assert output.count("error[") == 1, output
+        assert "type-assertion-failure" in output, output
+        assert "assert_never(puzzle_type)" in output, output
+        assert "FUTURE_KIND" in output, output
+
+
 # --- non-ASCII words end to end ----------------------------------------------