From 689c86a14cf9338674cc89a380d1cbe7f8c52fec Mon Sep 17 00:00:00 2001 From: Jusong Yu Date: Fri, 31 Jan 2025 21:48:07 +0100 Subject: [PATCH] More natural async flow By moving `self.future().result()` to step_until_terminated and await it in the event loop. It is more natural async control flow. --- src/plumpy/processes.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/plumpy/processes.py b/src/plumpy/processes.py index 94129464..c509706c 100644 --- a/src/plumpy/processes.py +++ b/src/plumpy/processes.py @@ -1314,10 +1314,10 @@ def execute(self) -> dict[str, Any] | None: :return: None if not terminated, otherwise `self.outputs` """ - if not self.has_terminated(): - asyncio.run(self.step_until_terminated()) + if self.has_terminated(): + return self.result() - return self.future().result() + return asyncio.run(self.step_until_terminated()) @ensure_not_closed async def step(self) -> None: @@ -1373,7 +1373,7 @@ async def step(self) -> None: self._stepping = False self._set_interrupt_action(None) - async def step_until_terminated(self) -> None: + async def step_until_terminated(self) -> Any: """If the process has not terminated, run the current step and wait until the step finished. @@ -1383,6 +1383,8 @@ async def step_until_terminated(self) -> None: while not self.has_terminated(): await self.step() + return await self.future() + # endregion @ensure_not_closed