refactor: replace AnyFunction by proper types (#1916)

Author: Kludex

src/mcp/server/fastmcp/server.py

@@ -6,7 +6,7 @@
 import re
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
-from typing import Any, Generic, Literal, overload
+from typing import Any, Generic, Literal, TypeVar, overload
 
 import anyio
 import pydantic_core
@@ -44,7 +44,7 @@
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared.context import LifespanContextT, RequestContext, RequestT
-from mcp.types import Annotations, AnyFunction, ContentBlock, GetPromptResult, Icon, ToolAnnotations
+from mcp.types import Annotations, ContentBlock, GetPromptResult, Icon, ToolAnnotations
 from mcp.types import Prompt as MCPPrompt
 from mcp.types import PromptArgument as MCPPromptArgument
 from mcp.types import Resource as MCPResource
@@ -53,6 +53,8 @@
 
 logger = get_logger(__name__)
 
+_CallableT = TypeVar("_CallableT", bound=Callable[..., Any])
+
 
 class Settings(BaseSettings, Generic[LifespanResultT]):
     """FastMCP server settings.
@@ -361,7 +363,7 @@ async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContent
 
     def add_tool(
         self,
-        fn: AnyFunction,
+        fn: Callable[..., Any],
         name: str | None = None,
         title: str | None = None,
         description: str | None = None,
@@ -417,7 +419,7 @@ def tool(
         icons: list[Icon] | None = None,
         meta: dict[str, Any] | None = None,
         structured_output: bool | None = None,
-    ) -> Callable[[AnyFunction], AnyFunction]:
+    ) -> Callable[[_CallableT], _CallableT]:
         """Decorator to register a tool.
 
         Tools can optionally request a Context object by adding a parameter with the
@@ -455,7 +457,7 @@ async def async_tool(x: int, context: Context) -> str:
                 "The @tool decorator was used incorrectly. Did you forget to call it? Use @tool() instead of @tool"
             )
 
-        def decorator(fn: AnyFunction) -> AnyFunction:
+        def decorator(fn: _CallableT) -> _CallableT:
             self.add_tool(
                 fn,
                 name=name,
@@ -507,7 +509,7 @@ def resource(
         icons: list[Icon] | None = None,
         annotations: Annotations | None = None,
         meta: dict[str, Any] | None = None,
-    ) -> Callable[[AnyFunction], AnyFunction]:
+    ) -> Callable[[_CallableT], _CallableT]:
         """Decorator to register a function as a resource.
 
         The function will be called when the resource is read to generate its content.
@@ -553,7 +555,7 @@ async def get_weather(city: str) -> str:
                 "Did you forget to call it? Use @resource('uri') instead of @resource"
             )
 
-        def decorator(fn: AnyFunction) -> AnyFunction:
+        def decorator(fn: _CallableT) -> _CallableT:
             # Check if this should be a template
             sig = inspect.signature(fn)
             has_uri_params = "{" in uri and "}" in uri
@@ -618,7 +620,7 @@ def prompt(
         title: str | None = None,
         description: str | None = None,
         icons: list[Icon] | None = None,
-    ) -> Callable[[AnyFunction], AnyFunction]:
+    ) -> Callable[[_CallableT], _CallableT]:
         """Decorator to register a prompt.
 
         Args:
@@ -660,7 +662,7 @@ async def analyze_file(path: str) -> list[Message]:
                 "Did you forget to call it? Use @prompt() instead of @prompt"
             )
 
-        def decorator(func: AnyFunction) -> AnyFunction:
+        def decorator(func: _CallableT) -> _CallableT:
             prompt = Prompt.from_function(func, name=name, title=title, description=description, icons=icons)
             self.add_prompt(prompt)
             return func
@@ -1086,7 +1088,7 @@ async def elicit(
 
         Args:
             schema: A Pydantic model class defining the expected response structure, according to the specification,
-                    only primive types are allowed.
+                    only primitive types are allowed.
             message: Optional message to present to the user. If not provided, will use
                     a default message based on the schema
 
@@ -1158,10 +1160,7 @@ async def log(
         """
 
         if extra:
-            log_data = {
-                "message": message,
-                **extra,
-            }
+            log_data = {"message": message, **extra}
         else:
             log_data = message
 

src/mcp/types.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Callable
 from datetime import datetime
 from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar
 
@@ -21,7 +20,6 @@
 Cursor = str
 Role = Literal["user", "assistant"]
 RequestId = Annotated[int, Field(strict=True)] | str
-AnyFunction: TypeAlias = Callable[..., Any]
 
 TaskExecutionMode = Literal["forbidden", "optional", "required"]
 TASK_FORBIDDEN: Final[Literal["forbidden"]] = "forbidden"