Propagate the return type

[dkang-quora]

Test code:

import asyncio
from typing import Any, Generator

from asynq import asynq


@asynq()
def noop():
    return


# case: no annotation + no yield
@asynq()
def f1():
    return 100


# case: no annotation + yield
@asynq()
def f2():
    yield noop.asynq()
    return 100


# case: annotation + no yield
@asynq()
def f3() -> int:
    return 300


# case: annotation + yield
@asynq()
def f4() -> Generator[Any, Any, int]:
    yield noop()
    return 300


async def main():
    _ret1: int = await f1.asyncio()
    _ret2: str = await f1.asyncio()  # "int" is incompatible with "str" (reportAssignmentType)

    _ret1: int = await f2.asyncio()
    _ret2: str = await f2.asyncio()  # "Literal[100]" is incompatible with "str" (reportAssignmentType)

    _ret1: int = await f3.asyncio()
    _ret2: str = await f3.asyncio()  # "int" is incompatible with "str" (reportAssignmentType)

    _ret1: int = await f4.asyncio()
    _ret2: str = await f4.asyncio()  # "int" is incompatible with "str" (reportAssignmentType)


asyncio.run(main())

• if we add asynq.G = Generator[Any, Any, T] for convenience and have pyanalyze accept both T and G[T], the existing asynq codebase can be converted to fully compatible code with pyright.

@asynq()
def f() -> G[int]:
  yield op.asynq()
  return 100

• Using TypeAlias(_G, _Coroutine, _CoroutineFn) did not work with pyright. I expanded them.
• The PR can wrongly unwrap Generator[, , T] as an edge case as I shared in the previous PR, but it's better than nothing

[JelleZijlstra]

Could you add a test case to the repo like in #144? These types are tricky enough that I'd like CI coverage to make sure it doesn't break. Or maybe we can merge #144 first and add this on top of that change.

[JelleZijlstra]

Internal testing shows that our current types break when methods are involved, I'll push something to try to address that.

[JelleZijlstra]

I removed the attempts to save the return type for now; they appear to be causing problems with generic functions. I'll merge this first to unblock internal use cases that are currently blocked, then try again to get return types to propagate.