[BUG] (some) type annotations for session wrong

[entropymatters]

Describe the bug A comment in #536 gives an example of using database sessions and transactions with beanie. However, when one wants to factor out the session into a separate function and annotates it, it does not properly type check using mypy anymore. Interestingly, this is not the case for all Document methods, despite beanie annotating session as pymongo.client_session.ClientSession everywhere.

I'd be glad to create a PR if the issue is really only to replace pymongo.client_session.ClientSession with motor.motor_asyncio.AsyncIOMotorClientSession, and would have a look in any other hints you can give for a fix.

Note that just casting the session to a pymongo ClientSession does not work since its start_session method is no coroutine.

To Reproduce

import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from motor import motor_asyncio
from beanie import Document, init_beanie

class Sample(Document):
    name: str

@asynccontextmanager
async def get_db_session(client: motor_asyncio.AsyncIOMotorClient) -> AsyncGenerator[motor_asyncio.AsyncIOMotorClientSession, None]:
    async with await client.start_session() as session:
        yield session

async def main() -> None:
    client = motor_asyncio.AsyncIOMotorClient('mongodb://mongo')
    database = client['mydatabase']
    await init_beanie(database=database, document_models=[Sample])

    async with get_db_session(client) as session:
        async with session.start_transaction():
            await Sample(name='alice').insert(session=session)  # OK
            await Sample.find(session=session).to_list()        # not OK

if __name__ == '__main__':
    asyncio.run(main())

Current behavior mypy shows an error for find() (but not for insert()):

error: No overload variant of "find" of "FindInterface" matches argument type "AsyncIOMotorClientSession"  [call-overload]
note: Possible overload variants:
note:     def find(cls, *args: Mapping[str, Any] | bool, projection_model: None = ..., skip: int | None = ..., limit: int | None = ..., sort: str | list[tuple[str, SortDirection]] | None = ..., session: ClientSession | None = ..., ignore_cache: bool = ..., fetch_links: bool = ..., with_children: bool = ..., lazy_parse: bool = ..., nesting_depth: int | None = ..., nesting_depths_per_field: dict[str, int] | None = ..., **pymongo_kwargs: Any) -> FindMany[Document]
note:     def [DocumentProjectionType <: BaseModel] find(cls, *args: Mapping[str, Any] | bool, projection_model: type[DocumentProjectionType], skip: int | None = ..., limit: int | None = ..., sort: str | list[tuple[str, SortDirection]] | None = ..., session: ClientSession | None = ..., ignore_cache: bool = ..., fetch_links: bool = ..., with_children: bool = ..., lazy_parse: bool = ..., nesting_depth: int | None = ..., nesting_depths_per_field: dict[str, int] | None = ..., **pymongo_kwargs: Any) -> FindMany[DocumentProjectionType]

Expected behavior Code type checks using mypy

Additional context none

[roman-right]

Hi! Thank you for the report. I'll pick it up during the next bug fixing session

[QSHolzner]

That would be nice. This is basically the same problem that I reported some time ago but unfortunately it was closed. see Issue: #706

[grthr]

I can confirm that this only is a typing issue. Passing the motor.motor_asyncio.AsyncIOMotorClientSession to the session parameter works fine and the transaction is acutally used. I worked around the typing issue with # type: ignore

async with await db.client.start_session() as session:
    async with session.start_transaction():
        await mydocument.insert(session=session)  # type: ignore