備忘録として。

やりたいこと

Discord.pyを使ったBotのプログラムに、以下があったとする。

@client.event
async def on_message(msg):
    if msg.content.startswith("counter"):
        count = 0
        while(True):
            print(count)
            count += 1
            await asyncio.sleep(1)

counterと送信されてからの秒数を一秒ごとにprintするプログラム。
while(True)でループしている上、中にretuenとかbreakとかする文もない。
なので、一度実行されたら延々と数字をカウントし続ける。これを強制終了させたい。

どうすればよいか

Discord.pyが使っているasyncioモジュールは、コルーチンが実行される度にtaskを形成する。
なので、関数内にasyncio.current_task()を記述することで実行されているコルーチンから、自分自身のtaskオブジェクトを取得できる。
その返り値をグローバル変数に代入し、それを参照してtask.cancel()する。

コードを書くと

stopと送信することで動きを止めたい場合以下の様になる:

#import asyncio をしておく
count_task = None
@client.event
async def on_message(msg):
    global count_task
    if msg.content.startswith("counter"):
        count = 0
        count_task = asyncio.current_task()
        while(True):
            print(count)
            count += 1
            await asyncio.sleep(1)
    elif msg.content.startswith("stop"):
        count_task.cancel()

最後に

これはあくまでも最低限の処理であり、counterやstopが連続で送信されたり(Botが入っている複数のサーバーで送信されるなど)すると例外が発生する。
複数のtaskを扱いたい場合は、taskオブジェクトを辞書に入れて管理するなどの工夫が必要。
また、正常にキャンセルされた場合でも、asyncio.CancelledErrorが出るので、それをエラーハンドラで拾う処理も書くと良い。