Python: async with の使い方

python 3.5 から使える async with の使い方です。
PEP 492 -- Coroutines with async and await syntax

async with?

コンテキストマネージャの async 版です。

1. DB にコネクションを張る <- 時間がかかる = async
2. 張ったコネクションでなにかする
3. コネクションを閉じる

みたいな処理があるときに

async with MyDB() as db:
    db.do_something()

というコードを実効するとして MyDB の方をうまく作っておけば async with ブロックを抜けたときに自動で db のコネクションを閉じるとかができます。

基本的な使い方

例えば、 aiomysql をラップしたクラスを作って、使ってみましょう。

my_db.py

my_db.py

import asyncio
import aiomysql


class MyDB:
    # async with に入る直前に呼ばれます。
    async def __aenter__(self):
        loop = asyncio.get_event_loop()
        self._connection = await aiomysql.connect(
            host='localhost',
            port=3306,
            user='root',
            password='ultrastrongpassword',
            db='my_db',
            charset='utf8',
            loop=loop
        )
        return self

    # async with ブロックを抜けた直後に呼ばれます。
    async def __aexit__(self, exc_type, exc, tb):
        self._connection.close()

    async def fetchall(self, query, args=[]):
        cursor = await self._connection.cursor()
        await cursor.execute(query, args)
        return await cursor.fetchall()

aiomysql の接続先とかを隠蔽した感じのものです。

大事なのは __aenter__ と __aexit__ です。
これらを実装することで、このクラスは非同期のコンテキストマネージャとして使えるようになります。
呼ばれるタイミングはコメントのとおりです。

例のように、 aenter でリソースのセットアップ、 aexit でリソースの開放をするのがよくあるパターンでしょう。

使う側

呼び出し側も見てみましょう。

使用する側

from my_db import MyDB


class Hoge:
    async def call_db(self):
        # db = MyDB()
        # async with db:
        # でも同じ意味だよ。
        async with MyDB() as db:
            result = db.fetchall('select * from some_table')
        print(result)

async with [非同期のコンテキストマネージャインスタンス] as [__aenter__ の返り値]: の書式になります。

これを実行すると、以下のような順番で処理が進みます。

1. MyDB のインスタンスが作られる（__init__ が呼ばれる）
2. Hoge の async with 節の処理が始まる
3. MyDB の __aenter__ が呼ばれる
4. MyDB の __aenter__ の返り値が async with as db の db に格納される
5. Hoge の result = db.fetchall('select * from some_table') が実行される
6. MyDB の __aexit__ が呼ばれる

async with の中にいるときだけ db インスタンスはもろもろのリソースを使えるようになる、という感じですね。

その他の使い方いろいろ

ファクトリで使いたい

async with MyDB.connect() as db: みたいにしたいとき

使用する側

class Hoge:
    async def call_db(self):
        async with MyDB.connect() as db:
            result = db.fetchall('select * from some_table')
        print(result)

以下のように普通に connect メソッドを実装するだけです。

my_db.py

class MyDB:
    __slots__ = ['_connection']

    async def __aenter__(self):
        (略)

    async def __aexit__(self, exc_type, exc, tb):
        self._connection.close()

    async def fetchall(self, query, args=[]):
        (略)

    @classmethod
    def connect(cls):
        return cls()

async with の後に MyDB （非同期コンテキストマネージャ）のインスタンスさえあれば良いので。

自分自身のインスタンスを返すんじゃない系で async with したい

使用する側

class Hoge:
    async def call_db(self):
        async with MyDB.direct_fetch('select * from some_table') as rows:
            print(rows)

この場合 async with の中身の MyDB.direct_fetch の返り値は MyDB インスタンスでは無さそうなので、ちょっと考える必要があります。
aiohttp の get とか が参考になるかもしれません。

以下のようにプライベートなコンテキストマネージャクラスを作ることで対応できます。

my_db.py

class _QueryContextManager:
    def __init__(self, coro):
        self._coro = coro

    async def __aenter__(self):
        loop = asyncio.get_event_loop()
        self._connection = await aiomysql.connect(
            host='localhost',
            port=3306,
            user='root',
            password='ultrastrongpassword',
            db='my_db',
            charset='utf8',
            loop=loop
        )
        return await self._coro(self._connection)

    async def __aexit__(self, exc_type, exc, tb):
        self._connection.close()


class MyDB:
    __slots__ = ['_connection']

    @classmethod
    def direct_fetch(cls, query, args=[]):
        async def fetch(connection):
            cursor = await connection.cursor()
            await cursor.execute(query, args)
            return await cursor.fetchall()

        # プライベートなコンテキストマネージャを返す。
        return _QueryContextManager(coro=fetch)

うーん。無理矢理感は否めませんね。。
もっと良い書き方はあるかもしれません。

aiohttp では request の返り値のインスタンスからコネクションを閉じれる（self._resp.release()）のでもう少しきれいになっているようです。