MAINT: P311-01 extend typing.Self to SeedDataset and Scorable factories

pyrit/models/score/scorable.py

@@ -5,7 +5,7 @@
 
 import uuid  # noqa: TC003  (runtime-required by Pydantic field annotations)
 from abc import ABC
-from typing import TYPE_CHECKING, Annotated, Any, Literal, get_args
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Self, get_args
 
 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator
 
@@ -66,15 +66,15 @@ def _validate_message_piece_ids(self) -> MessageScorable:
     def from_message(
         cls,
         message: Message,
-    ) -> MessageScorable:
+    ) -> Self:
         """
         Name the pieces of a persisted message.
 
         Args:
             message (Message): The message whose pieces to name.
 
         Returns:
-            MessageScorable: A scorable naming the message's pieces.
+            Self: A scorable naming the message's pieces.
         """
         return cls(message_piece_ids=tuple(piece.id for piece in message.message_pieces))
 
@@ -93,7 +93,7 @@ class ContentScorable(Scorable):
     data_type: PromptDataType = "text"
 
     @classmethod
-    def from_message(cls, message: Message) -> ContentScorable:
+    def from_message(cls, message: Message) -> Self:
         """
         Describe the converted content of a single-piece ephemeral message.
 
@@ -107,7 +107,7 @@ def from_message(cls, message: Message) -> ContentScorable:
             message (Message): The ephemeral message whose converted content to take.
 
         Returns:
-            ContentScorable: A scorable holding the converted message content.
+            Self: A scorable holding the converted message content.
         """
         piece = message.get_piece()
         return cls(value=piece.converted_value, data_type=piece.converted_value_data_type)

pyrit/models/seeds/seed_dataset.py

@@ -12,7 +12,7 @@
 import uuid
 from collections import defaultdict
 from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Self
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
@@ -233,7 +233,7 @@ def get_random_values(self, *, number: PositiveInt, harm_categories: Sequence[st
         return random.sample(prompts, min(len(prompts), number))
 
     @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> SeedDataset:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """
         Build a SeedDataset, assigning per-seed ``prompt_group_id`` by alias.
 
@@ -246,7 +246,7 @@ def from_dict(cls, data: dict[str, Any]) -> SeedDataset:
             data (dict[str, Any]): Dataset payload with top-level defaults and seed entries.
 
         Returns:
-            SeedDataset: Constructed dataset.
+            Self: Constructed dataset.
 
         Raises:
             ValueError: If any seed entry includes a pre-set ``prompt_group_id``.

tests/unit/models/test_scorable.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT license.
 
 import uuid
+from typing import assert_type
 
 import pytest
 from pydantic import ValidationError
@@ -80,6 +81,18 @@ def test_message_scorable_from_message_names_pieces():
     assert not hasattr(scorable, "message")
 
 
+def test_message_scorable_from_message_preserves_subclass_type():
+    class CustomMessageScorable(MessageScorable):
+        pass
+
+    message = _message()
+    scorable = CustomMessageScorable.from_message(message)
+
+    assert_type(scorable, CustomMessageScorable)
+    assert type(scorable) is CustomMessageScorable
+    assert scorable.message_piece_ids == (message.get_piece().id,)
+
+
 def test_message_scorable_rejects_empty_ids():
     with pytest.raises(ValueError, match="at least one message piece"):
         MessageScorable(message_piece_ids=())
@@ -118,6 +131,26 @@ def test_content_scorable_from_message_uses_converted_view():
     assert scorable.data_type == "text"
 
 
+def test_content_scorable_from_message_preserves_subclass_type():
+    class CustomContentScorable(ContentScorable):
+        pass
+
+    message = MessagePiece(
+        role="user",
+        original_value="original",
+        converted_value="converted",
+        original_value_data_type="text",
+        converted_value_data_type="text",
+    ).to_message()
+
+    scorable = CustomContentScorable.from_message(message)
+
+    assert_type(scorable, CustomContentScorable)
+    assert type(scorable) is CustomContentScorable
+    assert scorable.value == "converted"
+    assert scorable.data_type == "text"
+
+
 @pytest.mark.parametrize(
     "scorable",
     [

tests/unit/models/test_seed_dataset.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT license.
 
 import uuid
+from typing import assert_type
 
 import pytest
 
@@ -195,3 +196,14 @@ def test_seed_dataset_group_without_group_id():
 def test_seed_dataset_group_rejects_unsupported_seed_type():
     with pytest.raises(ValueError, match="Unsupported seed type: Seed"):
         SeedDataset.group_seed_prompts_by_prompt_group_id([Seed(value="unsupported")])
+
+
+def test_seed_dataset_from_dict_preserves_subclass_type():
+    class CustomSeedDataset(SeedDataset):
+        pass
+
+    ds = CustomSeedDataset.from_dict({"seeds": [{"value": "hello"}]})
+
+    assert_type(ds, CustomSeedDataset)
+    assert type(ds) is CustomSeedDataset
+    assert isinstance(ds.seeds[0], SeedPrompt)