diff --git a/src/async_kernel/caller.py b/src/async_kernel/caller.py index 37b1dfe6..f7d27e13 100644 --- a/src/async_kernel/caller.py +++ b/src/async_kernel/caller.py @@ -1042,13 +1042,11 @@ async def scheduler(): pen_.result() finally: queue.stop() - pen_.cancel() for pen in unfinished: pen.remove_done_callback(queue.append) if cancel_unfinished: pen.cancel("Cancelled by as_completed") - with anyio.CancelScope(): - await pen_.wait(result=False) + await pen_.cancel_wait(shield=True) async def wait( self, diff --git a/src/async_kernel/pending.py b/src/async_kernel/pending.py index 1eca9549..f858b092 100644 --- a/src/async_kernel/pending.py +++ b/src/async_kernel/pending.py @@ -612,12 +612,18 @@ def cancel(self, msg: str | None = None) -> bool: canceller(msg) return self._cancelled is not None - async def cancel_wait(self, msg: str | None, *, timeout: float | None = None) -> None: - "Cancel the pending and wait for it to be done." + async def cancel_wait(self, msg: str | None = None, *, timeout: float | None = None, shield: bool = False) -> None: + """ + Cancel the pending and wait for it to be done. + + Args: + timeout: Timeout in seconds. + shield: Shield from external cancellation. + """ if not self._done: self.cancel(msg) if not self._done: - await self.wait(result=False, timeout=timeout) + await self.wait(result=False, timeout=timeout, shield=shield) def cancelled(self) -> bool: """