Resolve postponed annotations before validation and schema creation

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -5,6 +5,7 @@ from fastapi import APIRouter, FastAPI, HTTPException, status
 from pydantic import BaseModel, create_model
 
 from ...helper import get_function_metadata_store
+from ...helper.get_function_type_hints import get_function_type_hints
 from ...views import Trace
 
 
@@ -47,9 +48,10 @@ def bootstrap_prediction_endpoint(
 
 def _get_schema(func: Callable) -> Type[BaseModel]:
     signature = inspect.signature(func)
+    type_hints = get_function_type_hints(func)
     parameters = {
         p.name: (
-            p.annotation if p.annotation != inspect._empty else Any,
+            type_hints.get(p.name, Any),
             p.default if p.default != inspect._empty else ...,
         )
         for p in signature.parameters.values()

great_ai/helper/get_function_type_hints.py

@@ -0,0 +1,7 @@
+import inspect
+from typing import Any, Callable, Dict, get_type_hints
+
+
+def get_function_type_hints(func: Callable) -> Dict[str, Any]:
+    """Resolve postponed annotations in the original function's namespace."""
+    return get_type_hints(inspect.unwrap(func), include_extras=True)

great_ai/parameters/parameter.py

@@ -6,6 +6,7 @@ from typeguard import TypeCheckError, check_type
 from ..errors import ArgumentValidationError
 from ..helper import get_arguments, get_function_metadata_store
 from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
+from ..helper.get_function_type_hints import get_function_type_hints
 from ..tracing.tracing_context import TracingContext
 
 F = TypeVar("F", bound=Callable)
@@ -56,14 +57,13 @@ def parameter(
         assert_function_is_not_finalised(func)
 
         actual_name = f"arg:{parameter_name}"
+        expected_type = get_function_type_hints(func).get(parameter_name)
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
             arguments = get_arguments(func, args, kwargs)
             argument = arguments.get(parameter_name)
 
-            expected_type = func.__annotations__.get(parameter_name)
-
             if expected_type is not None:
                 try:
                     check_type(argument, expected_type)

tests/test_argument_regressions.py

@@ -1,6 +1,8 @@
+import sys
+from types import ModuleType
+
 import pytest
 from fastapi.testclient import TestClient
-
 from great_ai import GreatAI
 
 
@@ -32,3 +34,25 @@ def test_positional_only_inputs_keep_their_bindings(isolated_context: None) -> N
         response = client.post("/predict", json={"first": "b", "second": 3})
         assert response.status_code == 200
         assert response.json()["output"] == "bbb"
+
+
+def test_postponed_model_annotations_resolve_in_their_module(
+    isolated_context: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    module = ModuleType("great_ai_test_annotations")
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    exec(
+        "from __future__ import annotations\n"
+        "from pydantic import BaseModel\n"
+        "class Payload(BaseModel):\n    value: int\n"
+        "def predict(data: Payload) -> int:\n    return data.value\n",
+        module.__dict__,
+    )
+    service = GreatAI.create(module.predict)
+    with pytest.raises(TypeError):
+        service("wrong")
+    with TestClient(service.app) as client:
+        response = client.post("/predict", json={"data": {"value": 7}})
+        assert response.status_code == 200
+        assert response.json()["output"] == 7
+        assert client.get("/openapi.json").status_code == 200