python - Asynchronous context manager

46
votes

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?

3
pythonasynchronouscontextmanager
asyncio has also introduced an async with asynchronous context manager protocol, see: python.org/dev/peps/pep-0492/… - jonrsharpe
This look like exaclty what I want. Will give a shot implementing it when I get a chance. - freebie
As of 3.7 (release somewhere in 2018), contextlib will have @asynccontextmanager - Yet Another User

3 Answers

63
votes

In Python 3.7, you'll be able to 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()

Until 3.7 comes out, 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()
30
votes

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.

9
votes

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