Skip to content

Commit

Permalink
More natural async flow
Browse files Browse the repository at this point in the history
By moving `self.future().result()` to step_until_terminated and await it
in the event loop. It is more natural async control flow.
  • Loading branch information
unkcpz committed Jan 31, 2025
1 parent 0901260 commit 689c86a
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions src/plumpy/processes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down

0 comments on commit 689c86a

Please sign in to comment.