Only wrap necessary parts of type aliases in forward annotations (#2548)

Author: Viicos

pydantic_ai_slim/pydantic_ai/tools.py

@@ -31,7 +31,7 @@
 ToolParams = ParamSpec('ToolParams', default=...)
 """Retrieval function param spec."""
 
-SystemPromptFunc = Union[
+SystemPromptFunc: TypeAlias = Union[
     Callable[[RunContext[AgentDepsT]], str],
     Callable[[RunContext[AgentDepsT]], Awaitable[str]],
     Callable[[], str],
@@ -42,25 +42,25 @@
 Usage `SystemPromptFunc[AgentDepsT]`.
 """
 
-ToolFuncContext = Callable[Concatenate[RunContext[AgentDepsT], ToolParams], Any]
+ToolFuncContext: TypeAlias = Callable[Concatenate[RunContext[AgentDepsT], ToolParams], Any]
 """A tool function that takes `RunContext` as the first argument.
 
 Usage `ToolContextFunc[AgentDepsT, ToolParams]`.
 """
-ToolFuncPlain = Callable[ToolParams, Any]
+ToolFuncPlain: TypeAlias = Callable[ToolParams, Any]
 """A tool function that does not take `RunContext` as the first argument.
 
 Usage `ToolPlainFunc[ToolParams]`.
 """
-ToolFuncEither = Union[ToolFuncContext[AgentDepsT, ToolParams], ToolFuncPlain[ToolParams]]
+ToolFuncEither: TypeAlias = Union[ToolFuncContext[AgentDepsT, ToolParams], ToolFuncPlain[ToolParams]]
 """Either kind of tool function.
 
 This is just a union of [`ToolFuncContext`][pydantic_ai.tools.ToolFuncContext] and
 [`ToolFuncPlain`][pydantic_ai.tools.ToolFuncPlain].
 
 Usage `ToolFuncEither[AgentDepsT, ToolParams]`.
 """
-ToolPrepareFunc: TypeAlias = 'Callable[[RunContext[AgentDepsT], ToolDefinition], Awaitable[ToolDefinition | None]]'
+ToolPrepareFunc: TypeAlias = Callable[[RunContext[AgentDepsT], 'ToolDefinition'], Awaitable['ToolDefinition | None']]
 """Definition of a function that can prepare a tool definition at call time.
 
 See [tool docs](../tools.md#tool-prepare) for more information.
@@ -88,9 +88,9 @@ def hitchhiker(ctx: RunContext[int], answer: str) -> str:
 Usage `ToolPrepareFunc[AgentDepsT]`.
 """
 
-ToolsPrepareFunc: TypeAlias = (
-    'Callable[[RunContext[AgentDepsT], list[ToolDefinition]], Awaitable[list[ToolDefinition] | None]]'
-)
+ToolsPrepareFunc: TypeAlias = Callable[
+    [RunContext[AgentDepsT], list['ToolDefinition']], Awaitable['list[ToolDefinition] | None']
+]
 """Definition of a function that can prepare the tool definition of all tools for each step.
 This is useful if you want to customize the definition of multiple tools or you want to register
 a subset of tools for a given step.
@@ -118,7 +118,7 @@ async def turn_on_strict_if_openai(
 Usage `ToolsPrepareFunc[AgentDepsT]`.
 """
 
-DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
+DocstringFormat: TypeAlias = Literal['google', 'numpy', 'sphinx', 'auto']
 """Supported docstring formats.
 
 * `'google'` — [Google-style](https://google.github.io/styleguide/pyguide.html#381-docstrings) docstrings.