Treat typing.Dict[str, X] return annotations like dict[str, X]

Author: Kludex

src/mcp/server/mcpserver/utilities/func_metadata.py

@@ -4,7 +4,6 @@
 import sys
 from collections.abc import Awaitable, Callable, Sequence
 from itertools import chain
-from types import GenericAlias
 from typing import Annotated, Any, NotRequired, Union, cast, get_args, get_origin
 
 import anyio
@@ -499,10 +498,8 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
         model = _create_wrapped_model(func_name, original_annotation)
         wrap_output = True
 
-    # Handle GenericAlias types (list[str], dict[str, int], Union[str, int], etc.)
-    elif isinstance(type_expr, GenericAlias):
-        origin = get_origin(type_expr)
-
+    # Handle parameterized generics (list[str], typing.Dict[str, int], Union[str, int], etc.)
+    elif (origin := get_origin(type_expr)) is not None:
         if origin is dict:
             args = get_args(type_expr)
             if len(args) == 2 and args[0] is str:
@@ -545,8 +542,6 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
 
     # Handle any other types not covered above
     else:
-        # This includes typing constructs that aren't GenericAlias in Python 3.10
-        # (e.g., Union, Optional in some Python versions)
         model = _create_wrapped_model(func_name, original_annotation)
         wrap_output = True
 

tests/server/mcpserver/test_func_metadata.py

@@ -3,6 +3,7 @@
 # pyright: reportMissingParameterType=false
 # pyright: reportUnknownArgumentType=false
 # pyright: reportUnknownLambdaType=false
+import typing
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, NotRequired, Required, TypedDict
@@ -295,6 +296,15 @@ def func_dict_list() -> dict[str, list[int]]:  # pragma: no cover
         "title": "func_dict_listDictOutput",
     }
 
+    def func_typing_dict() -> typing.Dict[str, str]:  # noqa: UP006  # pragma: no cover
+        return {"name": "John"}
+
+    assert func_metadata(func_typing_dict).output_schema == {
+        "type": "object",
+        "additionalProperties": {"type": "string"},
+        "title": "func_typing_dictDictOutput",
+    }
+
     # Test dict[int, str] - should be wrapped since key is not str
     def func_dict_int_key() -> dict[int, str]:  # pragma: no cover
         return {1: "a", 2: "b"}
@@ -774,6 +784,19 @@ def func_optional() -> str | None:  # pragma: no cover
         "title": "func_optionalOutput",
     }
 
+    UserId = typing.NewType("UserId", int)
+
+    def func_new_type() -> UserId:  # pragma: no cover
+        return UserId(1)
+
+    meta = func_metadata(func_new_type)
+    assert meta.output_schema == {
+        "type": "object",
+        "properties": {"result": {"title": "Result", "type": "integer"}},
+        "required": ["result"],
+        "title": "func_new_typeOutput",
+    }
+
 
 def test_structured_output_dataclass():
     """Test structured output with dataclass return types"""