MAINT: Modernize Optional/Union typing to PEP 604 syntax

pyrit/executor/attack/component/conversation_manager.py

@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from __future__ import annotations
+
 import logging
 import uuid
-from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from pyrit.common.deprecation import print_deprecation_message
 from pyrit.common.utils import combine_dict
@@ -22,14 +23,16 @@
     MessagePiece,
     Score,
 )
-from pyrit.prompt_normalizer.prompt_converter_configuration import (
-    PromptConverterConfiguration,
-)
 from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer
 from pyrit.prompt_target import CapabilityName, PromptTarget
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from pyrit.executor.attack.core import AttackContext
+    from pyrit.prompt_normalizer.prompt_converter_configuration import (
+        PromptConverterConfiguration,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -280,11 +283,11 @@ def set_system_prompt(
     async def initialize_context_async(
         self,
         *,
-        context: "AttackContext[Any]",
+        context: AttackContext[Any],
         target: PromptTarget,
         conversation_id: str,
         request_converters: list[PromptConverterConfiguration] | None = None,
-        prepended_conversation_config: Optional["PrependedConversationConfig"] = None,
+        prepended_conversation_config: PrependedConversationConfig | None = None,
         max_turns: int | None = None,
         memory_labels: dict[str, str] | None = None,
     ) -> ConversationState:
@@ -362,9 +365,9 @@ async def initialize_context_async(
     async def _handle_non_chat_target_async(
         self,
         *,
-        context: "AttackContext[Any]",
+        context: AttackContext[Any],
         prepended_conversation: list[Message],
-        config: Optional["PrependedConversationConfig"],
+        config: PrependedConversationConfig | None,
     ) -> ConversationState:
         """
         Handle prepended conversation for non-chat targets.
@@ -435,7 +438,7 @@ async def add_prepended_conversation_to_memory_async(
         prepended_conversation: list[Message],
         conversation_id: str,
         request_converters: list[PromptConverterConfiguration] | None = None,
-        prepended_conversation_config: Optional["PrependedConversationConfig"] = None,
+        prepended_conversation_config: PrependedConversationConfig | None = None,
         max_turns: int | None = None,
         target_identifier: ComponentIdentifier | None = None,
     ) -> int:
@@ -518,11 +521,11 @@ async def add_prepended_conversation_to_memory_async(
     async def _process_prepended_for_chat_target_async(
         self,
         *,
-        context: "AttackContext[Any]",
+        context: AttackContext[Any],
         prepended_conversation: list[Message],
         conversation_id: str,
         request_converters: list[PromptConverterConfiguration] | None,
-        prepended_conversation_config: Optional["PrependedConversationConfig"],
+        prepended_conversation_config: PrependedConversationConfig | None,
         max_turns: int | None,
         target_identifier: ComponentIdentifier | None = None,
     ) -> ConversationState:

pyrit/executor/attack/multi_turn/multi_prompt_sending.py

@@ -1,9 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from __future__ import annotations
+
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults
 from pyrit.common.utils import get_kwarg_param
@@ -51,13 +53,13 @@ class MultiPromptSendingAttackParameters(AttackParameters):
 
     @classmethod
     async def from_seed_group_async(
-        cls: type["MultiPromptSendingAttackParameters"],
+        cls: type[MultiPromptSendingAttackParameters],
         seed_group: SeedAttackGroup,
         *,
-        adversarial_chat: Optional["PromptTarget"] = None,
-        objective_scorer: Optional["TrueFalseScorer"] = None,
+        adversarial_chat: PromptTarget | None = None,
+        objective_scorer: TrueFalseScorer | None = None,
         **overrides: Any,
-    ) -> "MultiPromptSendingAttackParameters":
+    ) -> MultiPromptSendingAttackParameters:
         """
         Create parameters from a SeedGroup, extracting user messages.
 

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -1,14 +1,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from __future__ import annotations
+
 import asyncio
 import enum
 import json
 import logging
 import uuid
 from dataclasses import dataclass, field
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, cast, overload
+from typing import TYPE_CHECKING, Any, cast, overload
 
 from treelib.tree import Tree
 
@@ -70,6 +71,8 @@
 from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from pyrit.models.literals import PromptDataType
 
 logger = logging.getLogger(__name__)
@@ -171,7 +174,7 @@ class TAPAttackContext(MultiTurnAttackContext[Any]):
 
     # Nodes in the attack tree
     # Each node represents a branch in the attack tree with its own state
-    nodes: list["_TreeOfAttacksNode"] = field(default_factory=list)
+    nodes: list[_TreeOfAttacksNode] = field(default_factory=list)
 
     # Best conversation ID and score found during the attack
     best_conversation_id: str | None = None
@@ -376,7 +379,7 @@ async def initialize_with_prepended_conversation_async(
         self,
         *,
         prepended_conversation: list[Message],
-        prepended_conversation_config: Optional["PrependedConversationConfig"] = None,
+        prepended_conversation_config: PrependedConversationConfig | None = None,
     ) -> None:
         """
         Initialize the node with a prepended conversation history.
@@ -769,7 +772,7 @@ def _handle_unexpected_error(self, error: Exception) -> None:
         logger.error(f"Node {self.node_id}: Unexpected error during execution: {error}")
         self.error_message = f"Execution error: {str(error)}"
 
-    def duplicate(self) -> "_TreeOfAttacksNode":
+    def duplicate(self) -> _TreeOfAttacksNode:
         """
         Create a duplicate of this node for branching.
 

pyrit/memory/storage/serializers.py

@@ -60,7 +60,7 @@ def data_serializer_factory(
     Args:
         data_type (str): The type of the data (e.g., 'text', 'image_path', 'audio_path').
         value (str): The data value to be serialized.
-        extension (Optional[str]): The file extension, if applicable.
+        extension (str | None): The file extension, if applicable.
         category (AllowedCategories): The category or context for the data (e.g., 'seed-prompt-entries').
 
     Returns:
@@ -320,10 +320,10 @@ async def get_data_filename_async(self, file_name: str | None = None) -> Path |
         Generate or retrieve a unique filename for the data file.
 
         Args:
-            file_name (Optional[str]): Optional file name override.
+            file_name (str | None): Optional file name override.
 
         Returns:
-            Union[Path, str]: Full storage path for the generated data file.
+            Path | str: Full storage path for the generated data file.
 
         Raises:
             TypeError: If the serializer is not configured for on-disk data.
@@ -471,7 +471,7 @@ async def get_data_filename(  # pyrit-async-suffix-exempt
             file_name: Optional file name override.
 
         Returns:
-            Union[Path, str]: Full storage path for the generated data file.
+            Path | str: Full storage path for the generated data file.
         """
         print_deprecation_message(
             old_item="pyrit.memory.storage.serializers.DataTypeSerializer.get_data_filename",
@@ -586,7 +586,7 @@ def __init__(self, *, category: str, prompt_text: str, extension: str | None = N
         Args:
             category (str): Data category folder name.
             prompt_text (str): URL or path value.
-            extension (Optional[str]): Optional extension for persisted content.
+            extension (str | None): Optional extension for persisted content.
 
         """
         self.data_type = "url"
@@ -615,8 +615,8 @@ def __init__(self, *, category: str, prompt_text: str | None = None, extension:
 
         Args:
             category (str): Data category folder name.
-            prompt_text (Optional[str]): Optional existing image path.
-            extension (Optional[str]): Optional image extension.
+            prompt_text (str | None): Optional existing image path.
+            extension (str | None): Optional image extension.
 
         """
         self.data_type = "image_path"
@@ -652,8 +652,8 @@ def __init__(
 
         Args:
             category (str): Data category folder name.
-            prompt_text (Optional[str]): Optional existing audio path.
-            extension (Optional[str]): Optional audio extension.
+            prompt_text (str | None): Optional existing audio path.
+            extension (str | None): Optional audio extension.
 
         """
         self.data_type = "audio_path"
@@ -689,8 +689,8 @@ def __init__(
 
         Args:
             category (str): The category or context for the data.
-            prompt_text (Optional[str]): The video path or identifier.
-            extension (Optional[str]): The file extension, defaults to 'mp4'.
+            prompt_text (str | None): The video path or identifier.
+            extension (str | None): The file extension, defaults to 'mp4'.
 
         """
         self.data_type = "video_path"
@@ -730,8 +730,8 @@ def __init__(
 
         Args:
             category (str): The category or context for the data.
-            prompt_text (Optional[str]): The binary file path or identifier.
-            extension (Optional[str]): The file extension, defaults to 'bin'.
+            prompt_text (str | None): The binary file path or identifier.
+            extension (str | None): The file extension, defaults to 'bin'.
 
         """
         self.data_type = "binary_path"

pyrit/memory/storage/storage.py

@@ -70,7 +70,7 @@ async def read_file(self, path: Path | str) -> bytes:  # pyrit-async-suffix-exem
         Read a file from storage (deprecated alias of ``read_file_async``).
 
         Args:
-            path (Union[Path, str]): The path to the file.
+            path (Path | str): The path to the file.
 
         Returns:
             bytes: The content of the file.
@@ -87,7 +87,7 @@ async def write_file(self, path: Path | str, data: bytes) -> None:  # pyrit-asyn
         Write data to storage (deprecated alias of ``write_file_async``).
 
         Args:
-            path (Union[Path, str]): The path to the file.
+            path (Path | str): The path to the file.
             data (bytes): The content to write to the file.
         """
         print_deprecation_message(
@@ -102,7 +102,7 @@ async def path_exists(self, path: Path | str) -> bool:  # pyrit-async-suffix-exe
         Check whether a path exists (deprecated alias of ``path_exists_async``).
 
         Args:
-            path (Union[Path, str]): The path to check.
+            path (Path | str): The path to check.
 
         Returns:
             bool: True if the path exists, False otherwise.
@@ -119,7 +119,7 @@ async def is_file(self, path: Path | str) -> bool:  # pyrit-async-suffix-exempt
         Check whether the given path is a file (deprecated alias of ``is_file_async``).
 
         Args:
-            path (Union[Path, str]): The path to check.
+            path (Path | str): The path to check.
 
         Returns:
             bool: True if the path is a file, False otherwise.
@@ -136,7 +136,7 @@ async def create_directory_if_not_exists(self, path: Path | str) -> None:  # pyr
         Create a directory if it does not exist (deprecated alias of ``create_directory_if_not_exists_async``).
 
         Args:
-            path (Union[Path, str]): The directory path to create.
+            path (Path | str): The directory path to create.
         """
         print_deprecation_message(
             old_item="pyrit.memory.storage.storage.StorageIO.create_directory_if_not_exists",
@@ -156,7 +156,7 @@ async def read_file_async(self, path: Path | str) -> bytes:
         Asynchronously reads a file from the local disk.
 
         Args:
-            path (Union[Path, str]): The path to the file.
+            path (Path | str): The path to the file.
 
         Returns:
             bytes: The content of the file.
@@ -224,7 +224,7 @@ def _convert_to_path(self, path: Path | str) -> Path:
         Convert an input path to a Path object.
 
         Args:
-            path (Union[Path, str]): Input path value.
+            path (Path | str): Input path value.
 
         Returns:
             Path: Normalized Path instance.
@@ -249,8 +249,8 @@ def __init__(
         Initialize an Azure Blob Storage I/O adapter.
 
         Args:
-            container_url (Optional[str]): Azure Blob container URL.
-            sas_token (Optional[str]): Optional SAS token.
+            container_url (str | None): Azure Blob container URL.
+            sas_token (str | None): Optional SAS token.
             blob_content_type (SupportedContentType): Blob content type for uploads.
 
         Raises:
@@ -363,7 +363,7 @@ def _resolve_blob_name(self, path: Path | str) -> str:
         created on Windows still produce valid blob names.
 
         Args:
-            path (Union[Path, str]): Blob URL or relative blob path.
+            path (Path | str): Blob URL or relative blob path.
 
         Returns:
             str: The resolved blob name.
@@ -428,7 +428,7 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None:
         If a relative path is provided, it is used as the blob name directly.
 
         Args:
-            path (Union[Path, str]): Full blob URL or relative blob path.
+            path (Path | str): Full blob URL or relative blob path.
             data (bytes): The data to write.
         """
         if not self._client_async:
@@ -448,7 +448,7 @@ async def path_exists_async(self, path: Path | str) -> bool:
         Check whether a given path exists in the Azure Blob Storage container.
 
         Args:
-            path (Union[Path, str]): Blob URL or path to test.
+            path (Path | str): Blob URL or path to test.
 
         Returns:
             bool: True when the path exists.
@@ -473,7 +473,7 @@ async def is_file_async(self, path: Path | str) -> bool:
         Check whether the path refers to a file (blob) in Azure Blob Storage.
 
         Args:
-            path (Union[Path, str]): Blob URL or path to test.
+            path (Path | str): Blob URL or path to test.
 
         Returns:
             bool: True when the blob exists and has non-zero content size.
@@ -498,7 +498,7 @@ async def create_directory_if_not_exists_async(self, directory_path: Path | str)
         Log a no-op directory creation for Azure Blob Storage.
 
         Args:
-            directory_path (Union[Path, str]): Requested directory path.
+            directory_path (Path | str): Requested directory path.
 
         """
         logger.info(

pyrit/prompt_converter/prompt_converter.py

@@ -1,12 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from __future__ import annotations
+
 import abc
 import asyncio
 import inspect
 import re
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, get_args
+from typing import TYPE_CHECKING, Any, ClassVar, get_args
 
 from pyrit import prompt_converter
 from pyrit.models import ComponentIdentifier, ConverterIdentifier, Identifiable, PromptDataType
@@ -89,7 +91,7 @@ def __init_subclass__(cls, **kwargs: object) -> None:
                     f"Declare the output modalities this converter produces."
                 )
 
-    def __init__(self, *, converter_target: Optional["PromptTarget"] = None) -> None:
+    def __init__(self, *, converter_target: PromptTarget | None = None) -> None:
         """
         Initialize the prompt converter.