sqlmodel Fix AsyncSession annotations

Fixes #54

`hero.py`:

import asyncio
import inspect
import sys
from contextlib import suppress
from pathlib import Path
from typing import Optional, TYPE_CHECKING

from mypy.main import main as mypy_main
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import Field, SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession


class Hero(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str
    age: Optional[int] = None


async def main() -> None:
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")

    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)

    async with AsyncSession(engine) as session:
        session.add(Hero(name="Spider-Boy", secret_name="Pedro Parqueador"))
        await session.commit()

    async with AsyncSession(engine) as session:
        statement = select(Hero).where(Hero.name == "Spider-Boy")
        reveal_type(session)
        result = await session.exec(statement)
        reveal_type(result)
        h = result.first()
        reveal_type(h)


if not TYPE_CHECKING:

    def reveal_type(obj):
        f_back = inspect.currentframe().f_back
        filename = f_back.f_globals["__file__"]
        with suppress(ValueError):
            filename = Path(filename).relative_to(Path.cwd())
        print(f"{filename}:{f_back.f_lineno}: note: Runtime value is {obj!r}")


print("Running main()")
asyncio.run(main())
print("\nRunning mypy (first run may take some time)")
mypy_main(None, stdout=sys.stdout, stderr=sys.stderr, args=[__file__])

Running on main

Running main()
hero.py:33: note: Runtime value is <sqlmodel.ext.asyncio.session.AsyncSession object at 0x10314f9a0>
hero.py:35: note: Runtime value is <sqlalchemy.engine.result.ScalarResult object at 0x10308bdc0>
hero.py:37: note: Runtime value is Hero(age=None, name='Spider-Boy', secret_name='Pedro Parqueador', id=1)

Running mypy (first run may take some time)
hero.py:33: note: Revealed type is 'sqlmodel.ext.asyncio.session.AsyncSession*'
hero.py:34: error: Need type annotation for 'result'
hero.py:34: error: Argument 1 to "exec" of "AsyncSession" has incompatible type "SelectOfScalar[Hero]"; expected "Union[Select[<nothing>], Executable[<nothing>]]"
hero.py:35: note: Revealed type is 'sqlmodel.engine.result.ScalarResult[Any]'
hero.py:37: note: Revealed type is 'Union[Any, None]'
Found 2 errors in 1 file (checked 1 source file)

Running on Bobronium:AsyncSession_typing_fix:

Running main()
hero.py:43: note: Runtime value is <sqlmodel.ext.asyncio.session.AsyncSession object at 0x102a0f940>
hero.py:45: note: Runtime value is <sqlalchemy.engine.result.ScalarResult object at 0x1029e9190>
hero.py:47: note: Runtime value is Hero(name='Spider-Boy', secret_name='Pedro Parqueador', id=1, age=None)

Running mypy (first run may take some time)
hero.py:43: note: Revealed type is 'sqlmodel.ext.asyncio.session.AsyncSession*'
hero.py:45: note: Revealed type is 'sqlmodel.engine.result.ScalarResult*[hero.Hero*]'
hero.py:47: note: Revealed type is 'Union[hero.Hero*, None]'

Aug 29 '21 10:08 Bobronium

@tiangolo, is this PR or #54 going to be addressed?

Jun 30 '22 17:06 Bobronium

Thanks for the feature! Until this PR gets merged, manually patching will work for my workflow.

Jul 17 '22 11:07 priyansh-anand

📝 Docs preview for commit 11ff44ff6206362370460830c286db0aa51d9325 at: https://639cfc052b35ad15e2e48091--sqlmodel.netlify.app

Dec 16 '22 23:12 github-actions[bot]

@tiangolo, I can try to resolve CI issues and update the branch, but I need a confirmation that you have intention of merging it once issues are resolved and willing to communicate if I'll need any feedback on resolving them. Complete lack of comminucation from your side on this issue is frustrating.

I don't want to waste any more time on this PR only to find it in the same condition a year later.

Jan 12 '23 17:01 Bobronium

Thanks @Bobronium for the fix, I am pip installing your branch instead of v0.0.8. I hope this gets merged soon.