Asynchronous context manager

Question

I have an asynchronous API which I'm using to connect and send mail to an SMTP server which has some setup and tear down to it. So it fits nicely into using a contextmanager from Python 3's contextlib.

Though, I don't know if it's possible write because they both use the generator syntax to write.

This might demonstrate the problem (contains a mix of yield-base and async-await syntax to demonstrate the difference between async calls and yields to the context manager).

@contextmanager
async def smtp_connection():
    client = SMTPAsync()
    ...

    try:
        await client.connect(smtp_url, smtp_port)
        await client.starttls()
        await client.login(smtp_username, smtp_password)
        yield client
    finally:
        await client.quit()

Is this kind of thing possible within python currently? and how would I use a with as statement if it is? If not is there a alternative way I could achieve this - maybe using the old style context manager?

`asyncio` has also introduced an `async with` asynchronous context manager protocol, see: https://www.python.org/dev/peps/pep-0492/#asynchronous-context-managers-and-async-with — jonrsharpe, May 25 '16 at 09:42
This look like exaclty what I want. Will give a shot implementing it when I get a chance. — freebie, May 25 '16 at 14:08
As of 3.7 (release somewhere in 2018), contextlib will have `@asynccontextmanager` — Yet Another User, Jun 25 '17 at 17:11
related: https://stackoverflow.com/questions/3693771/understanding-the-python-with-statement-and-context-managers — Charlie Parker, Jul 29 '22 at 21:09
how do you know what functions to call in the try and finally when using @asynccontextmanager? e.g. my custom class has an `__aexit__` and `__aenter__`. Do I call them manually myself? — Charlie Parker, Aug 03 '22 at 17:29

score 83 · Accepted Answer · edited Aug 21 '21 at 19:06

Since Python 3.7, you can write:

from contextlib import asynccontextmanager

@asynccontextmanager
async def smtp_connection():
    client = SMTPAsync()
    ...

    try:
        await client.connect(smtp_url, smtp_port)
        await client.starttls()
        await client.login(smtp_username, smtp_password)
        yield client
    finally:
        await client.quit()

Before 3.7, you can use the async_generator package for this. On 3.6, you can write:

# This import changed, everything else is the same
from async_generator import asynccontextmanager

@asynccontextmanager
async def smtp_connection():
    client = SMTPAsync()
    ...

    try:
        await client.connect(smtp_url, smtp_port)
        await client.starttls()
        await client.login(smtp_username, smtp_password)
        yield client
    finally:
        await client.quit()

And if you want to work all the way back to 3.5, you can write:

# This import changed again:
from async_generator import asynccontextmanager, async_generator, yield_

@asynccontextmanager
@async_generator      # <-- added this
async def smtp_connection():
    client = SMTPAsync()
    ...

    try:
        await client.connect(smtp_url, smtp_port)
        await client.starttls()
        await client.login(smtp_username, smtp_password)
        await yield_(client)    # <-- this line changed
    finally:
        await client.quit()

how do you know what functions to call in the try and finally? e.g. my custom class has an `__aexit__` and `__aenter__`. Do I call them manually myself? — Charlie Parker, Aug 03 '22 at 17:27
I think it's important to clarify from the docs or mention that your decorated function must `The function being decorated must return a generator-iterator when called.` otherwise it's confusing what one needs to do: https://docs.python.org/3/library/contextlib.html — Charlie Parker, Aug 03 '22 at 18:17

score 43 · Answer 2 · answered May 25 '16 at 15:42

43

Thanks to @jonrsharpe was able to make an async context manager.

Here's what mine ended up looking like for anyone who want's some example code:

class SMTPConnection():
    def __init__(self, url, port, username, password):
        self.client   = SMTPAsync()
        self.url      = url
        self.port     = port
        self.username = username
        self.password = password

    async def __aenter__(self):
        await self.client.connect(self.url, self.port)
        await self.client.starttls()
        await self.client.login(self.username, self.password)

        return self.client

    async def __aexit__(self, exc_type, exc, tb):
        await self.client.quit()

usage:

async with SMTPConnection(url, port, username, password) as client:
    await client.sendmail(...)

Feel free to point out if I've done anything stupid.

answered May 25 '16 at 15:42

freebie

2,161
2
19
36

2

The problem is that if you use it twice simultaneously, your second enter's client will overwrite the first one and the exit on first one will then exit the second too. – iScrE4m Feb 12 '18 at 13:56
@iScrE4m Ah yes, I hadn't anticipated it being used multiple times, just creating a instance, on-demand, for one-off use. Could maybe make a wrapper class to this which delegates it's `__aenter__` and `__aexit__` to new instances. – freebie Feb 12 '18 at 14:49
4

Remember to handle exceptions in your `__aexit__`. Otherwise you'll hide exceptions and get weird invisible bugs. – Thomas Ahle Aug 23 '18 at 11:16
AIU this constitutes "doing it manually", and `asynccontextmanager` is the modern way to achieve this, so this answer should be marked as obsolete maybe? – P i Aug 21 '21 at 19:09
how do you know what functions to call in the try and finally? e.g. my custom class has an `__aexit__` and `__aenter__`. Do I call them manually myself? – Charlie Parker Aug 03 '22 at 17:27
or does everyone have a `.close()`/`.quite()` method? – Charlie Parker Aug 03 '22 at 17:37

score 8 · Answer 3 · edited Feb 17 '20 at 23:03

8

The asyncio_extras package has a nice solution for this:

import asyncio_extras

@asyncio_extras.async_contextmanager
async def smtp_connection():
    client = SMTPAsync()
    ...

For Python < 3.6, you'd also need the async_generator package and replace yield client with await yield_(client).

edited Feb 17 '20 at 23:03

Andrii Tykhonov

538
6
11

answered Jan 04 '17 at 18:27

Bart Robinson

1,024
1
9
10

how do you know what functions to call in the try and finally? e.g. my custom class has an `__aexit__` and `__aenter__`. Do I call them manually myself? – Charlie Parker Aug 03 '22 at 17:27
or does everyone have a `.close()`/`.quite()` method? – Charlie Parker Aug 03 '22 at 17:37
I think it's important to clarify from the docs or mention that your decorated function must `The function being decorated must return a generator-iterator when called.` otherwise it's confusing what one needs to do: https://docs.python.org/3/library/contextlib.html – Charlie Parker Aug 03 '22 at 18:17

score 0 · Answer 4 · answered Aug 03 '22 at 19:58

I find that you need to call obj.__aenter__(...) in the try and obj.__aexit__(...) in the final. Perhaps you do too if all you want is abstract an overly complicated object that has resources.

e.g.

import asyncio
from contextlib import asynccontextmanager

from pycoq.common import CoqContext, LocalKernelConfig
from pycoq.serapi import CoqSerapi

from pdb import set_trace as st


@asynccontextmanager
async def get_coq_serapi(coq_ctxt: CoqContext) -> CoqSerapi:
    """
    Returns CoqSerapi instance that is closed with a with statement.
    CoqContext for the file is also return since it can be used to manipulate the coq file e.g. return
    the coq statements as in for `stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):`.

    example use:
    ```
    filenames = pycoq.opam.opam_strace_build(coq_package, coq_package_pin)
        filename: str
        for filename in filenames:
            with get_coq_serapi(filename) as coq, coq_ctxt:
                for stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):
    ```

    ref:
    - https://stackoverflow.com/questions/37433157/asynchronous-context-manager
    - https://stackoverflow.com/questions/3693771/understanding-the-python-with-statement-and-context-managers

    Details:

    Meant to replace (see Brando's pycoq tutorial):
    ```
            async with aiofile.AIOFile(filename, 'rb') as fin:
                coq_ctxt = pycoq.common.load_context(filename)
                cfg = opam.opam_serapi_cfg(coq_ctxt)
                logfname = pycoq.common.serapi_log_fname(os.path.join(coq_ctxt.pwd, coq_ctxt.target))
                async with pycoq.serapi.CoqSerapi(cfg, logfname=logfname) as coq:
    ```
    usually then you loop through the coq stmts e.g.
    ```
                    for stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):
    ```
    """
    try:
        import pycoq
        from pycoq import opam
        from pycoq.common import LocalKernelConfig
        import os

        # - note you can't return the coq_ctxt here so don't create it due to how context managers work, even if it's needed layer for e.g. stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):
        # _coq_ctxt: CoqContext = pycoq.common.load_context(coq_filepath)
        # - not returned since it seems its only needed to start the coq-serapi interface
        cfg: LocalKernelConfig = opam.opam_serapi_cfg(coq_ctxt)
        logfname = pycoq.common.serapi_log_fname(os.path.join(coq_ctxt.pwd, coq_ctxt.target))
        # - needed to be returned to talk to coq
        coq: CoqSerapi = pycoq.serapi.CoqSerapi(cfg, logfname=logfname)
        # - crucial, or coq._kernel is None and .execute won't work
        await coq.__aenter__()  # calls self.start(), this  must be called by itself in the with stmt beyond yield
        yield coq
    except Exception as e:
        # fin.close()
        # coq.close()
        import traceback
        await coq.__aexit__(Exception, e, traceback.format_exc())
        # coq_ctxt is just a data class serapio no need to close it, see: https://github.com/brando90/pycoq/blob/main/pycoq/common.py#L32
    finally:
        import traceback
        err_msg: str = 'Finally exception clause'
        exception_type, exception_value = Exception('Finally exception clause'), ValueError(err_msg)
        print(f'{traceback.format_exc()=}')
        await coq.__aexit__(exception_type, exception_value, traceback.format_exc())
        # coq_ctxt is just a data class so no need to close it, see: https://github.com/brando90/pycoq/blob/main/pycoq/common.py#L32


# -

async def loop_through_files_original():
    ''' '''
    import os

    import aiofile

    import pycoq
    from pycoq import opam

    coq_package = 'lf'
    from pycoq.test.test_autoagent import with_prefix
    coq_package_pin = f"file://{with_prefix('lf')}"

    print(f'{coq_package=}')
    print(f'{coq_package_pin=}')
    print(f'{coq_package_pin=}')

    filenames: list[str] = pycoq.opam.opam_strace_build(coq_package, coq_package_pin)
    filename: str
    for filename in filenames:
        print(f'-> {filename=}')
        async with aiofile.AIOFile(filename, 'rb') as fin:
            coq_ctxt: CoqContext = pycoq.common.load_context(filename)
            cfg: LocalKernelConfig = opam.opam_serapi_cfg(coq_ctxt)
            logfname = pycoq.common.serapi_log_fname(os.path.join(coq_ctxt.pwd, coq_ctxt.target))
            async with pycoq.serapi.CoqSerapi(cfg, logfname=logfname) as coq:
                print(f'{coq._kernel=}')
                for stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):
                    print(f'--> {stmt=}')
                    _, _, coq_exc, _ = await coq.execute(stmt)
                    if coq_exc:
                        raise Exception(coq_exc)


async def loop_through_files():
    """
    to test run in linux:
    ```
        python ~pycoq/pycoq/utils.py
        python -m pdb -c continue ~/pycoq/pycoq/utils.py
    ```
    """
    import pycoq

    coq_package = 'lf'
    from pycoq.test.test_autoagent import with_prefix
    coq_package_pin = f"file://{with_prefix('lf')}"

    print(f'{coq_package=}')
    print(f'{coq_package_pin=}')
    print(f'{coq_package_pin=}')

    filenames: list[str] = pycoq.opam.opam_strace_build(coq_package, coq_package_pin)
    filename: str
    for filename in filenames:
        print(f'-> {filename=}')
        coq_ctxt: CoqContext = pycoq.common.load_context(filename)
        async with get_coq_serapi(coq_ctxt) as coq:
            print(f'{coq=}')
            print(f'{coq._kernel=}')
            stmt: str
            for stmt in pycoq.split.coq_stmts_of_context(coq_ctxt):
                print(f'--> {stmt=}')
                _, _, coq_exc, _ = await coq.execute(stmt)
                if coq_exc:
                    raise Exception(coq_exc)


if __name__ == '__main__':
    asyncio.run(loop_through_files_original())
    asyncio.run(loop_through_files())
    print('Done!\a\n')

see code: https://github.com/brando90/pycoq/blob/main/pycoq/utils.py

Asynchronous context manager

4 Answers4

Linked

Related