🐛 Fix AsyncSession type annotations for exec()

[Bobronium]

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]'

[Bobronium]

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

[priyansh-anand]

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

[github-actions]

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

[Bobronium]

@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.

[maresb]

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

[diego-escobedo]

Crazy this has been around since 2021. Awesome work @Bobronium on this.

[ornakash]

Please merge this?

[Bobronium]

@tiangolo, I've fixed CI issue.

Tests are passing now.

Please merge.

[tiangolo]

Great, thanks @Bobronium! 🚀

Thanks everyone for the comments, in particular confirming this solved it for you. 🍰

Thanks for the patience everyone!

This will be available in the next version, 0.0.9. 🎉