Python: resolve string forward references nested inside list/dict/tuple annotations

python/semantic_kernel/schema/kernel_json_schema_builder.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import ast
+import builtins
 import sys
 import types
 from enum import Enum
@@ -86,6 +88,7 @@ def build_model_schema(
         hints = get_type_hints(model, globalns=model_module_globals, localns={})
 
         for field_name, field_type in hints.items():
+            field_type = cls._resolve_nested_forward_refs(field_type, model_module_globals)
             field_description = None
             if hasattr(model, "model_fields") and field_name in model.model_fields:
                 field_info = model.model_fields[field_name]
@@ -150,6 +153,141 @@ def get_json_schema(cls, parameter_type: type) -> dict[str, Any]:
         type_name = TYPE_MAPPING.get(parameter_type, "object")
         return {"type": type_name}
 
+    _TYPE_EXPRESSION_NODES = (
+        ast.Expression,
+        ast.Name,
+        ast.Attribute,
+        ast.Subscript,
+        ast.Tuple,
+        ast.List,
+        ast.Load,
+        ast.Constant,
+        ast.BinOp,
+        ast.BitOr,
+    )
+
+    @classmethod
+    def _is_type_expression(cls, source: str) -> bool:
+        """Return whether `source` parses as a type expression and nothing more.
+
+        A forward reference names a type: an identifier, a dotted path, a subscription such as
+        `dict[str, Inner]`, or a `X | None` union. Calls, lambdas and comprehensions are not part
+        of that grammar, so rejecting them limits `_evaluate_type_expression` to the few
+        operations a type annotation is made of.
+        """
+        try:
+            tree = ast.parse(source, mode="eval")
+        except SyntaxError:
+            return False
+        for node in ast.walk(tree):
+            if not isinstance(node, cls._TYPE_EXPRESSION_NODES):
+                return False
+            if isinstance(node, ast.Constant) and not isinstance(node.value, (str, int, bool, type(None))):
+                return False
+        return True
+
+    @classmethod
+    def _evaluate_type_expression(cls, node: ast.AST, globalns: dict[str, Any]) -> Any:
+        """Evaluate a tree that `_is_type_expression` admitted, without using `eval`.
+
+        Names are looked up the way `eval` looks them up, in `globalns` first and then in the
+        builtins. Every other admitted node is a subscription, an attribute access, a `|` union,
+        a tuple or list of arguments, or a literal; there is no call node to evaluate.
+        """
+        if isinstance(node, ast.Expression):
+            return cls._evaluate_type_expression(node.body, globalns)
+        if isinstance(node, ast.Name):
+            if node.id in globalns:
+                return globalns[node.id]
+            return getattr(builtins, node.id)
+        if isinstance(node, ast.Attribute):
+            return getattr(cls._evaluate_type_expression(node.value, globalns), node.attr)
+        if isinstance(node, ast.Subscript):
+            value = cls._evaluate_type_expression(node.value, globalns)
+            return value[cls._evaluate_type_expression(node.slice, globalns)]
+        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
+            left = cls._evaluate_type_expression(node.left, globalns)
+            return left | cls._evaluate_type_expression(node.right, globalns)
+        if isinstance(node, ast.Tuple):
+            return tuple(cls._evaluate_type_expression(element, globalns) for element in node.elts)
+        if isinstance(node, ast.List):
+            return [cls._evaluate_type_expression(element, globalns) for element in node.elts]
+        if isinstance(node, ast.Constant):
+            return node.value
+        raise TypeError(f"Unsupported node in a type expression: {type(node).__name__}")
+
+    @classmethod
+    def _is_type_or_typing_construct(cls, value: Any) -> bool:
+        """Return whether `value` is something a forward reference can name.
+
+        That is a class (a model, an enum, a builtin) or a typing construct such as
+        `dict[str, Inner]`, `Optional[Inner]` or `Inner | None`, all of which `get_origin` recognizes.
+        """
+        return isinstance(value, type) or get_origin(value) is not None
+
+    @classmethod
+    def _resolve_nested_forward_refs(cls, annotation: Any, globalns: dict[str, Any]) -> Any:
+        """Resolve string forward references nested inside a generic alias.
+
+        `get_type_hints` evaluates an annotation that *is* a string, but it does not descend into
+        a generic alias that already exists as an object. `list["Inner"]` goes through
+        `list.__class_getitem__`, which stores `"Inner"` verbatim instead of wrapping it in a
+        `ForwardRef`, so nothing resolves it and `build` formats the bare string rather than the
+        class it names.
+
+        Args:
+            annotation: The annotation to resolve, typically a generic alias.
+            globalns: The globals of the module the owning model was defined in.
+
+        Returns:
+            Any: The annotation with resolvable string arguments replaced by the types they name,
+                or the original annotation when a reference does not resolve to a type.
+        """
+        args = get_args(annotation)
+        if not args:
+            return annotation
+
+        resolved_args = []
+        changed = False
+        for arg in args:
+            reference = arg if isinstance(arg, str) else getattr(arg, "__forward_arg__", None)
+            if reference is not None:
+                if not cls._is_type_expression(reference):
+                    # Anything that isn't the grammar of a type expression is not a forward
+                    # reference; refuse to evaluate it rather than widen what a stray
+                    # annotation can run while a schema is being built.
+                    return annotation
+                try:
+                    resolved = cls._evaluate_type_expression(ast.parse(reference, mode="eval"), globalns)
+                except Exception:
+                    # Not resolvable from this module; leave the annotation alone so the existing
+                    # fallback applies instead of raising while a schema is being built.
+                    return annotation
+                if not cls._is_type_or_typing_construct(resolved):
+                    # The name exists but is a module, a function or a plain value rather than a
+                    # type. Substituting it would make `build` raise or describe the wrong object,
+                    # so keep the reference as written and let the existing fallback apply.
+                    return annotation
+                changed = True
+            else:
+                resolved = cls._resolve_nested_forward_refs(arg, globalns)
+                changed = changed or resolved is not arg
+            resolved_args.append(resolved)
+
+        if not changed:
+            return annotation
+
+        copy_with = getattr(annotation, "copy_with", None)
+        if copy_with is not None:
+            return copy_with(tuple(resolved_args))
+        origin = get_origin(annotation)
+        if origin is None:
+            return annotation
+        try:
+            return origin[tuple(resolved_args)]
+        except TypeError:
+            return annotation
+
     @classmethod
     def handle_complex_type(
         cls, parameter_type: type, description: str | None = None, structured_output: bool = False

python/tests/unit/schema/test_schema_builder.py

@@ -1,8 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import ast
 import json
+import sys
 from enum import Enum
-from typing import Annotated, Any, Optional, Union
+from typing import Annotated, Any, ForwardRef, Optional, Union, get_args
 from unittest.mock import Mock
 
 import pytest
@@ -455,3 +457,219 @@ def test_build_schema_with_nonpydantic_structured_output():
     }
 
     assert structured_output_schema == expected_schema
+
+
+# Up to 3.10, get_type_hints leaves a string nested in list[...] as written, so the builder is what resolves it.
+# From 3.11 on, get_type_hints evaluates that string itself, and the builder only ever sees the result.
+skip_if_get_type_hints_evaluates_nested_strings = pytest.mark.skipif(
+    sys.version_info >= (3, 11),
+    reason="From 3.11 on, get_type_hints evaluates strings nested in list[...] itself, before the builder runs.",
+)
+
+
+class InnerForward(KernelBaseModel):
+    value: int
+    label: str
+
+
+class HolderForwardList(KernelBaseModel):
+    items: list["InnerForward"] = []
+
+
+class HolderDirectList(KernelBaseModel):
+    items: list[InnerForward] = []
+
+
+class HolderForwardDict(KernelBaseModel):
+    mapping: dict[str, "InnerForward"] = {}
+
+
+class HolderForwardOptional(KernelBaseModel):
+    maybe: Optional["InnerForward"] = None
+
+
+def test_build_list_with_string_forward_reference_matches_direct_reference():
+    """`list["Inner"]` and `list[Inner]` must produce the same schema."""
+    forward = KernelJsonSchemaBuilder.build(HolderForwardList)
+    direct = KernelJsonSchemaBuilder.build(HolderDirectList)
+
+    assert forward == direct
+    assert forward["properties"]["items"]["items"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+def test_build_dict_with_string_forward_reference():
+    schema = KernelJsonSchemaBuilder.build(HolderForwardDict)
+
+    assert schema["properties"]["mapping"]["additionalProperties"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+def test_build_optional_with_string_forward_reference_still_works():
+    """`Optional["Inner"]` already resolved via `get_type_hints`; make sure it still does."""
+    schema = KernelJsonSchemaBuilder.build(HolderForwardOptional)
+
+    assert schema["properties"]["maybe"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+@skip_if_get_type_hints_evaluates_nested_strings
+def test_unresolvable_forward_reference_falls_back_instead_of_raising():
+    """An annotation naming something that doesn't exist must not break schema building."""
+
+    class HolderUnknown(KernelBaseModel):
+        items: list["DoesNotExistAnywhere"] = []  # noqa: F821
+
+    schema = KernelJsonSchemaBuilder.build(HolderUnknown)
+
+    assert schema["properties"]["items"]["type"] == "array"
+    assert schema["properties"]["items"]["items"] == {"type": "object"}
+
+
+def test_non_type_expression_forward_reference_is_not_evaluated():
+    """A nested string that isn't a type expression must not be evaluated.
+
+    `list["Inner"]` stores the string verbatim, so whatever it contains reaches the resolver.
+    Only the grammar of a type expression is evaluated; a call is left alone and the annotation
+    comes back unchanged for the existing fallback to handle.
+    """
+    executed = []
+
+    def _canary():
+        executed.append(True)
+        return int
+
+    annotation = list["_canary()"]  # noqa: F821
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, {"_canary": _canary})
+
+    assert executed == []
+    assert resolved is annotation
+
+
+def test_type_expression_forward_reference_is_still_resolved():
+    """The guard must not block the case the resolver exists for."""
+    annotation = list["InnerForward"]  # noqa: F821
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, {"InnerForward": InnerForward})
+
+    assert get_args(resolved) == (InnerForward,)
+
+
+@pytest.mark.parametrize(
+    ("reference", "is_type_expression"),
+    [
+        ("InnerForward", True),
+        ("dict[str, InnerForward]", True),
+        ("InnerForward | None", True),
+        ("tuple[int, str]", True),
+        ("__import__('os').getcwd()", False),
+        ("_canary()", False),
+        ("(lambda: 1)()", False),
+        ("[x for x in ().__class__.__base__.__subclasses__()]", False),
+    ],
+)
+def test_is_type_expression(reference: str, is_type_expression: bool):
+    """The guard admits type expressions and rejects anything that can call out."""
+    assert KernelJsonSchemaBuilder._is_type_expression(reference) is is_type_expression
+
+
+@pytest.mark.parametrize(
+    "reference",
+    [
+        "InnerForward",
+        "dict[str, InnerForward]",
+        "InnerForward | None",
+        "tuple[int, str]",
+        "Optional[InnerForward]",
+        "Union[int, str, None]",
+        "Annotated[int, 'meta']",
+        "list[dict[str, 'InnerForward']]",
+        "json.JSONDecoder",
+    ],
+)
+def test_evaluate_type_expression_matches_eval(reference: str):
+    """Every expression the guard admits evaluates to the same object `eval` would return."""
+    assert KernelJsonSchemaBuilder._is_type_expression(reference)
+    tree = ast.parse(reference, mode="eval")
+
+    assert KernelJsonSchemaBuilder._evaluate_type_expression(tree, globals()) == eval(reference, globals(), {})
+
+
+def test_unknown_name_in_forward_reference_is_left_unchanged():
+    annotation = list["DoesNotExistAnywhere"]  # noqa: F821
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, globals())
+
+    assert resolved is annotation
+
+
+def forward_reference_helper(count: int) -> str:
+    return str(count)
+
+
+FORWARD_REFERENCE_LIMITS = [1, 2, 3]
+FORWARD_REFERENCE_ALIAS = "int"
+
+
+class HolderModuleReference:
+    items: list["json"]
+
+
+class HolderFunctionReference:
+    items: list["forward_reference_helper"]
+
+
+class HolderListValueReference:
+    items: list["FORWARD_REFERENCE_LIMITS"]
+
+
+class HolderStrValueReference:
+    items: list["FORWARD_REFERENCE_ALIAS"]
+
+
+@pytest.mark.parametrize(
+    "reference", ["json", "forward_reference_helper", "FORWARD_REFERENCE_LIMITS", "FORWARD_REFERENCE_ALIAS"]
+)
+@pytest.mark.parametrize("wrap", [str, ForwardRef], ids=["str", "ForwardRef"])
+def test_forward_reference_to_non_type_global_is_left_unchanged(reference: str, wrap: Any):
+    """A name that resolves to a module, a function or a plain value does not name a type.
+
+    Substituting it would hand `build` a module (raises), a function (rendered as a model of its
+    parameters), a list (raises as an unhashable mapping key) or a string (read as another type
+    name). The annotation must come back unchanged so the existing fallback applies.
+    """
+    annotation = list[wrap(reference)]
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, globals())
+
+    assert resolved is annotation
+
+
+@pytest.mark.parametrize(
+    ("reference", "expected"),
+    [
+        ("InnerForward", InnerForward),
+        ("dict[str, InnerForward]", dict[str, InnerForward]),
+        ("Optional[InnerForward]", Optional[InnerForward]),
+        ("InnerForward | None", InnerForward | None),
+    ],
+)
+def test_forward_reference_to_type_or_typing_construct_is_resolved(reference: str, expected: Any):
+    """Classes and typing constructs are what a forward reference names; those are still substituted."""
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(list[reference], globals())
+
+    assert get_args(resolved) == (expected,)
+
+
+@skip_if_get_type_hints_evaluates_nested_strings
+@pytest.mark.parametrize(
+    "holder",
+    [HolderModuleReference, HolderFunctionReference, HolderListValueReference, HolderStrValueReference],
+)
+def test_build_with_non_type_forward_reference_keeps_object_fallback(holder: type):
+    schema = KernelJsonSchemaBuilder.build(holder)
+
+    assert schema["properties"]["items"] == {"type": "array", "items": {"type": "object"}}