From bbc6f78efab52bd72416299b3582c0b58bef06f7 Mon Sep 17 00:00:00 2001 From: Tai Sakuma Date: Tue, 11 Jun 2024 19:50:43 -0400 Subject: [PATCH] Set asyncio events in the finally block --- nextline/fsm/state.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/nextline/fsm/state.py b/nextline/fsm/state.py index 58343d8a..ae3cd529 100644 --- a/nextline/fsm/state.py +++ b/nextline/fsm/state.py @@ -45,11 +45,16 @@ async def on_enter_running(self, _: EventData) -> None: run_started = asyncio.Event() async def run() -> None: - async with self._hook.awith.run(context=self._context): - run_started.set() - self._context.run_arg = None - await self.finish() # type: ignore - self.run_finished.set() + try: + async with self._hook.awith.run(context=self._context): + run_started.set() + finally: + run_started.set() # Ensure to unblock the await + self._context.run_arg = None + try: + await self.finish() # type: ignore + finally: + self.run_finished.set() self._task = asyncio.create_task(run()) await run_started.wait() @@ -71,7 +76,7 @@ async def on_close_while_running(self, _: EventData) -> None: async def wait(self) -> None: await self.run_finished.wait() - + async def on_enter_finished(self, _: EventData) -> None: await self._hook.ahook.on_finished(context=self._context)