Incorrect typing when using asyncio.run

[Ben-Epstein]

Apologies if i'm in the wrong repo. I thought this was the right place because of #30

I have an example where mypy is complaining that my call is returning a type object when it should be returning Response[T] | Response[str]

import asyncio
import json
from typing import Generic, Type, TypeVar, overload

from pydantic import BaseModel, Field

T = TypeVar("T", bound=BaseModel)
U = TypeVar("U", bound=str | BaseModel) 


class Response(BaseModel, Generic[U]):
    text: str
    response_schema: Type[BaseModel] | None

    @property
    def out(self) -> U:
        if self.response_schema:
            return self.response_schema.model_validate_json(self.text)
        return self.text

    class Config:
        arbitrary_types_allowed = True


@overload
async def apredict(prompt: str, response_schema: Type[T]) -> Response[T]: ...
@overload
async def apredict(prompt: str, response_schema: None = ...) -> Response[str]: ...

async def apredict(prompt: str, response_schema: Type[T] | None = None) -> Response[str] | Response[T]:
    ... # irrelevant 
    
@overload
def predict(prompt: str, response_schema: Type[T]) -> Response[T]: ...
@overload
def predict(prompt: str, response_schema: None = ...) -> Response[str]: ...

def predict(prompt: str, response_schema: Type[T] | None = None) -> Response[str] | Response[T]:
    return asyncio.run(apredict(prompt, response_schema))

I get an error Incompatible return value type (got "object", expected "Response[str] | Response[T]") which I don't expect because it seems to be typed generically based on the signature of the function it is running.

thanks!

[srittau]

This is due to the narrowing behavior of mypy. See this simplified example:

import asyncio
from typing import overload

@overload
async def apredict(x: int) -> int: ...
@overload
async def apredict(x: str) -> str: ...
async def apredict(x: int | str) -> int | str:
    return ""
    
x: int | str
y = asyncio.run(apredict(x))
reveal_type(y)  # object

Unfortunately, there's nothing we can do about this in typeshed, although I think there's talk about changing this behavior over at mypy.