MAINT: Type adversarial placeholder factory with Self

pyrit/models/messages/message_piece.py

@@ -5,7 +5,7 @@
 
 import uuid
 from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Self
 from uuid import uuid4
 
 from pydantic import (
@@ -226,7 +226,7 @@ def is_truncated(self) -> bool:
     # Adversarial placeholder support
     # ------------------------------------------------------------------ #
     @classmethod
-    def adversarial_placeholder(cls, *, role: ChatMessageRole = "user") -> MessagePiece:
+    def adversarial_placeholder(cls, *, role: ChatMessageRole = "user") -> Self:
         """
         Build a placeholder text piece that signals the adversarial chat will
         generate the text content at this position.

tests/unit/models/test_message_piece.py

@@ -6,6 +6,7 @@
 import uuid
 from collections.abc import MutableSequence
 from datetime import UTC, datetime, timedelta
+from typing import assert_type
 from unittest.mock import patch
 
 import pytest
@@ -862,6 +863,19 @@ def test_message_piece_has_error_and_is_blocked_consistency():
     assert blocked_entry.has_error() is True
 
 
+class _CustomMessagePiece(MessagePiece):
+    """A subclass used to check classmethod factory return types."""
+
+
+def test_adversarial_placeholder_factory_returns_calling_subclass():
+    piece = _CustomMessagePiece.adversarial_placeholder(role="assistant")
+
+    assert_type(piece, _CustomMessagePiece)
+    assert type(piece) is _CustomMessagePiece
+    assert piece.role == "assistant"
+    assert piece.is_adversarial_placeholder() is True
+
+
 def test_adversarial_placeholder_factory_defaults():
     piece = MessagePiece.adversarial_placeholder()