ContextManager doesn't work with async, or maybe ExplicitContainer?

[jeffsawatzky]

Trying to use the context container to create a "transaction" for the whole "request". Here is some code to give an example:

import asyncio
import dataclasses
from collections.abc import AsyncIterator
from typing import Any

import lagom


class Transaction:
    async def __aenter__(self) -> "Transaction":
        return self

    async def commit(self) -> None:
        return

    async def rollback(self) -> None:
        return

    async def close(self) -> None:
        return

    async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any, **kwargs: Any) -> None:
        if exc_type is None:
            # Auto commit if there was no exception
            await self.commit()
        else:
            # Auto rollback is there was an exception
            await self.rollback()
        return None


@dataclasses.dataclass
class UserModel:
    user_id: int


class UserRepository:
    def __init__(self, transaction: Transaction) -> None:
        self._transaction = transaction

    async def find_user(self, user_id: int) -> UserModel:
        user = None
        async with self._transaction as t:
            # do some db stuff here, maybe create the user if it doesn't exist or something...
            user = UserModel(user_id=user_id)
        return user


class FindUserUseCase:
    def __init__(self, user_repo: UserRepository) -> None:
        self._user_repo = user_repo

    async def execute(self, user_id: int) -> UserModel:
        return await self._user_repo.find_user(user_id=user_id)


container = lagom.ExplicitContainer()


@lagom.context_dependency_definition(container)  # type: ignore[misc]
async def _get_a_transaction() -> AsyncIterator[Transaction]:
    transaction = Transaction()
    try:
        yield transaction
    finally:
        await transaction.close()


container[UserRepository] = lambda c: UserRepository(c[Transaction])

container[FindUserUseCase] = lambda c: FindUserUseCase(c[UserRepository])

context_container = lagom.ContextContainer(container=container, context_types=[], context_singletons=[Transaction])


@lagom.bind_to_container(context_container)
async def find_user(user_id: int, find_user_use_case: FindUserUseCase = lagom.injectable) -> None:
    user = await find_user_use_case.execute(user_id=user_id)
    print(user)


if __name__ == "__main__":
    asyncio.run(find_user(1))

When I run this I get:

lagom.exceptions.InvalidDependencyDefinition: A ContextManager[<class '__main__.Transaction'>] should be defined. This could be an Iterator[<class '__main__.Transaction'>] or Generator[<class '__main__.Transaction'>, None, None] with the @contextmanager decorator

Am I doing something wrong? If so, what?

[jeffsawatzky]

@meadsteve I got it working by looking at the AsyncContextManager stuff in the experimental module.

A couple things that tripped me up though:

1. You use typing.Awaitable but the recommended import is collections.abc.Awaitable. I was using the wrong Awaitable and it wasn't being found in the container.
2. The AwaitableSingleton seems out of place/very snowflakey. Currently I don't have a better approach though.
3. The bound function needs to specify the Awaitable type, and then await it to get the dependency. It would be nice if the bound function could get the actual dependency provided to it. In other words, have the bind_to_containerdecorator do the awaiting.
4. The [Async]ContextContainer takes a context_singletons list for types that are "shared" within a single "request", and the bind_to_container takes a shared list of types for the same purpose. Is there a way to consolidate these two?

import asyncio
import dataclasses
from typing import Any, AsyncIterator, Awaitable  # noqa

import lagom
from lagom.experimental.context_based import AsyncContextContainer, AwaitableSingleton


class Transaction:
    def __init__(self) -> None:
        pass

    async def __aenter__(self) -> "Transaction":
        return self

    async def commit(self) -> None:
        return

    async def rollback(self) -> None:
        return

    async def close(self) -> None:
        return

    async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any, **kwargs: Any) -> None:
        if exc_type is None:
            # Auto commit if there was no exception
            await self.commit()
        else:
            # Auto rollback is there was an exception
            await self.rollback()
        return None


@dataclasses.dataclass
class UserModel:
    user_id: int


class UserRepository:
    def __init__(self, transaction: Transaction) -> None:
        self._transaction = transaction

    async def find_user(self, user_id: int) -> UserModel:
        user = None
        async with self._transaction:
            # do some db stuff here, maybe create the user if it doesn't exist or something...
            user = UserModel(user_id=user_id)
        return user


class FindUserUseCase:
    def __init__(self, user_repo: UserRepository) -> None:
        self._user_repo = user_repo

    async def execute(self, user_id: int) -> UserModel:
        return await self._user_repo.find_user(user_id=user_id)


container = lagom.ExplicitContainer()


@lagom.context_dependency_definition(container)  # type: ignore[misc]
async def provide_transaction() -> AsyncIterator[Transaction]:
    transaction = Transaction()
    try:
        yield transaction
    finally:
        await transaction.close()


@lagom.dependency_definition(container)  # type: ignore[misc]
async def provide_user_repository(c: lagom.Container) -> UserRepository:
    transaction = await c[AwaitableSingleton[Transaction]].get()
    return UserRepository(transaction)


@lagom.dependency_definition(container)  # type: ignore[misc]
async def provide_find_user_use_case(c: lagom.Container) -> FindUserUseCase:
    user_repository = await c[Awaitable[UserRepository]]  # type:ignore[type-abstract]
    return FindUserUseCase(user_repository)


context_container = AsyncContextContainer(container=container, context_types=[], context_singletons=[Transaction])


@lagom.bind_to_container(context_container)
async def find_user(user_id: int, find_user_use_case_awaitable: Awaitable[FindUserUseCase] = lagom.injectable) -> None:
    find_user_use_case = await find_user_use_case_awaitable
    user = await find_user_use_case.execute(user_id=user_id)
    print(user)


if __name__ == "__main__":
    asyncio.run(find_user(1))