Skip to content

Commit d79c72e

Browse files
committed
add closing() context manager, which by default does a shutdown(wait=False) for better performance in async
1 parent 5a1838b commit d79c72e

File tree

3 files changed

+28
-1
lines changed

3 files changed

+28
-1
lines changed

examples/executor_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ async def master():
1616

1717
await first_50(progress)
1818
loop = asyncio.get_running_loop()
19-
with QThreadExecutor(1) as exec:
19+
with QThreadExecutor(1).closing() as exec:
2020
await loop.run_in_executor(exec, functools.partial(last_50, progress), loop)
2121

2222

src/qasync/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,14 @@ def __enter__(self, *args):
258258
def __exit__(self, *args):
259259
self.shutdown()
260260

261+
@contextlib.contextmanager
262+
def closing(self, *, wait=False, cancel_futures=False):
263+
"""Explicit context manager to do shutdown, with Wait=False by default"""
264+
try:
265+
yield self
266+
finally:
267+
self.shutdown(wait=wait, cancel_futures=cancel_futures)
268+
261269

262270
def _result_or_cancel(fut, timeout=None):
263271
try:

tests/test_qthreadexec.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import time
88
import weakref
99
from concurrent.futures import CancelledError, TimeoutError
10+
from unittest import mock
1011

1112
import pytest
1213

@@ -196,3 +197,21 @@ def func(x):
196197
# because the max number of workers is 5 and the rest of
197198
# the tasks were not started at the time of the cancel.
198199
assert set(results) != {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
200+
201+
202+
def test_closing(executor):
203+
"""Test that closing context manager works as expected"""
204+
# mock the shutdown method of the executor
205+
with mock.patch.object(executor, "shutdown") as mock_shutdown:
206+
with executor.closing():
207+
pass
208+
209+
# ensure that shutdown was called with (False, cancel_futures=False)
210+
mock_shutdown.assert_called_once_with(wait=False, cancel_futures=False)
211+
212+
with mock.patch.object(executor, "shutdown") as mock_shutdown:
213+
with executor.closing(wait=True, cancel_futures=True):
214+
pass
215+
216+
# ensure that shutdown was called with (False, cancel_futures=False)
217+
mock_shutdown.assert_called_once_with(wait=True, cancel_futures=True)

0 commit comments

Comments
 (0)