Skip to content

Commit c6807c5

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

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
@@ -8,6 +8,7 @@
88
import weakref
99
from concurrent.futures import CancelledError, TimeoutError
1010
from itertools import islice
11+
from unittest import mock
1112

1213
import pytest
1314

@@ -246,3 +247,21 @@ def func(x):
246247
m.close()
247248
executor.shutdown(wait=True, cancel_futures=False)
248249
assert len(results) < 10, "Some tasks should have been cancelled"
250+
251+
252+
def test_closing(executor):
253+
"""Test that closing context manager works as expected"""
254+
# mock the shutdown method of the executor
255+
with mock.patch.object(executor, "shutdown") as mock_shutdown:
256+
with executor.closing():
257+
pass
258+
259+
# ensure that shutdown was called with (False, cancel_futures=False)
260+
mock_shutdown.assert_called_once_with(wait=False, cancel_futures=False)
261+
262+
with mock.patch.object(executor, "shutdown") as mock_shutdown:
263+
with executor.closing(wait=True, cancel_futures=True):
264+
pass
265+
266+
# ensure that shutdown was called with (False, cancel_futures=False)
267+
mock_shutdown.assert_called_once_with(wait=True, cancel_futures=True)

0 commit comments

Comments
 (0)