Improve variance of classes (#726)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/_result.py

@@ -8,17 +8,20 @@
 from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin
 
 from pydantic import TypeAdapter, ValidationError
-from typing_extensions import Self, TypeAliasType, TypedDict
+from typing_extensions import TypeAliasType, TypedDict, TypeVar
 
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
-from .result import ResultData, ResultValidatorFunc
-from .tools import AgentDeps, RunContext, ToolDefinition
+from .result import ResultDataT, ResultDataT_inv, ResultValidatorFunc
+from .tools import AgentDepsT, RunContext, ToolDefinition
+
+T = TypeVar('T')
+"""An invariant TypeVar."""
 
 
 @dataclass
-class ResultValidator(Generic[AgentDeps, ResultData]):
-    function: ResultValidatorFunc[AgentDeps, ResultData]
+class ResultValidator(Generic[AgentDepsT, ResultDataT_inv]):
+    function: ResultValidatorFunc[AgentDepsT, ResultDataT_inv]
     _takes_ctx: bool = field(init=False)
     _is_async: bool = field(init=False)
 
@@ -28,10 +31,10 @@ def __post_init__(self):
 
     async def validate(
         self,
-        result: ResultData,
+        result: T,
         tool_call: _messages.ToolCallPart | None,
-        run_context: RunContext[AgentDeps],
-    ) -> ResultData:
+        run_context: RunContext[AgentDepsT],
+    ) -> T:
         """Validate a result but calling the function.
 
         Args:
@@ -50,10 +53,10 @@ async def validate(
 
         try:
             if self._is_async:
-                function = cast(Callable[[Any], Awaitable[ResultData]], self.function)
+                function = cast(Callable[[Any], Awaitable[T]], self.function)
                 result_data = await function(*args)
             else:
-                function = cast(Callable[[Any], ResultData], self.function)
+                function = cast(Callable[[Any], T], self.function)
                 result_data = await _utils.run_in_executor(function, *args)
         except ModelRetry as r:
             m = _messages.RetryPromptPart(content=r.message)
@@ -74,17 +77,19 @@ def __init__(self, tool_retry: _messages.RetryPromptPart):
 
 
 @dataclass
-class ResultSchema(Generic[ResultData]):
+class ResultSchema(Generic[ResultDataT]):
     """Model the final response from an agent run.
 
     Similar to `Tool` but for the final result of running an agent.
     """
 
-    tools: dict[str, ResultTool[ResultData]]
+    tools: dict[str, ResultTool[ResultDataT]]
     allow_text_result: bool
 
     @classmethod
-    def build(cls, response_type: type[ResultData], name: str, description: str | None) -> Self | None:
+    def build(
+        cls: type[ResultSchema[T]], response_type: type[T], name: str, description: str | None
+    ) -> ResultSchema[T] | None:
         """Build a ResultSchema dataclass from a response type."""
         if response_type is str:
             return None
@@ -95,10 +100,10 @@ def build(cls, response_type: type[ResultData], name: str, description: str | No
         else:
             allow_text_result = False
 
-        def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[ResultData]:
-            return cast(ResultTool[ResultData], ResultTool(a, tool_name_, description, multiple))
+        def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[T]:
+            return cast(ResultTool[T], ResultTool(a, tool_name_, description, multiple))
 
-        tools: dict[str, ResultTool[ResultData]] = {}
+        tools: dict[str, ResultTool[T]] = {}
         if args := get_union_args(response_type):
             for i, arg in enumerate(args, start=1):
                 tool_name = union_tool_name(name, arg)
@@ -112,7 +117,7 @@ def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[ResultDat
 
     def find_named_tool(
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
-    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultData]] | None:
+    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultDataT]] | None:
         """Find a tool that matches one of the calls, with a specific name."""
         for part in parts:
             if isinstance(part, _messages.ToolCallPart):
@@ -122,7 +127,7 @@ def find_named_tool(
     def find_tool(
         self,
         parts: Iterable[_messages.ModelResponsePart],
-    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultData]] | None:
+    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultDataT]] | None:
         """Find a tool that matches one of the calls."""
         for part in parts:
             if isinstance(part, _messages.ToolCallPart):
@@ -142,11 +147,11 @@ def tool_defs(self) -> list[ToolDefinition]:
 
 
 @dataclass(init=False)
-class ResultTool(Generic[ResultData]):
+class ResultTool(Generic[ResultDataT]):
     tool_def: ToolDefinition
     type_adapter: TypeAdapter[Any]
 
-    def __init__(self, response_type: type[ResultData], name: str, description: str | None, multiple: bool):
+    def __init__(self, response_type: type[ResultDataT], name: str, description: str | None, multiple: bool):
         """Build a ResultTool dataclass from a response type."""
         assert response_type is not str, 'ResultTool does not support str as a response type'
 
@@ -183,7 +188,7 @@ def __init__(self, response_type: type[ResultData], name: str, description: str
 
     def validate(
         self, tool_call: _messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
-    ) -> ResultData:
+    ) -> ResultDataT:
         """Validate a result message.
 
         Args:

pydantic_ai_slim/pydantic_ai/_system_prompt.py

@@ -6,12 +6,12 @@
 from typing import Any, Callable, Generic, cast
 
 from . import _utils
-from .tools import AgentDeps, RunContext, SystemPromptFunc
+from .tools import AgentDepsT, RunContext, SystemPromptFunc
 
 
 @dataclass
-class SystemPromptRunner(Generic[AgentDeps]):
-    function: SystemPromptFunc[AgentDeps]
+class SystemPromptRunner(Generic[AgentDepsT]):
+    function: SystemPromptFunc[AgentDepsT]
     dynamic: bool = False
     _takes_ctx: bool = field(init=False)
     _is_async: bool = field(init=False)
@@ -20,7 +20,7 @@ def __post_init__(self):
         self._takes_ctx = len(inspect.signature(self.function).parameters) > 0
         self._is_async = inspect.iscoroutinefunction(self.function)
 
-    async def run(self, run_context: RunContext[AgentDeps]) -> str:
+    async def run(self, run_context: RunContext[AgentDepsT]) -> str:
         if self._takes_ctx:
             args = (run_context,)
         else: