agronholm / asyncio_extras Goto Github PK

class TransactionAbort(IOError):
    pass

@async_contextmanager
async def transaction(conn):
    try:
        yield conn
    except TransactionAbort:
        ...
        return
    else:
        ...

        try:
            await self.generator.athrow(exc_val)
        except StopAsyncIteration:
            return True  # tell python that the exception (StopAsyncIteration) is safe to suppress

from asyncio_extras import contextmanager


def patch_context_manager():
    if hasattr(contextmanager._AsyncContextManager.__aexit__, '__horcruxed__'):
        return

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_val is not None:
            try:
                await self.generator.athrow(exc_val)
            except StopAsyncIteration:
                return True
        else:
            try:
                await self.generator.asend(None)
            except StopAsyncIteration:
                pass
            else:
                raise RuntimeError("async generator didn't stop")
    __aexit__.__horcruxed__ = time.time()
    contextmanager._AsyncContextManager.__aexit__ = __aexit__

@asyncio_extras.async_contextmanager
async def scoped_session():
    session = Session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()

async def create_api_key(module):
    async with scoped_session() as session:
        pass
    pass

$ pylint -E file.py
No config file found, using default configuration
************* Module file
E:  6,10: Async context manager 'generator' doesn't implement __aenter__ and __aexit__. (not-async-context-manager)