備忘録として。
やりたいこと
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が出るので、それをエラーハンドラで拾う処理も書くと良い。