add closing() context manager, which by default does a shutdown(wait=False) for better performance in async

Author: kristjanvalur

examples/executor_example.py

@@ -16,7 +16,7 @@ async def master():
 
     await first_50(progress)
     loop = asyncio.get_running_loop()
-    with QThreadExecutor(1) as exec:
+    with QThreadExecutor(1).closing() as exec:
         await loop.run_in_executor(exec, functools.partial(last_50, progress), loop)
 
 

src/qasync/__init__.py

@@ -258,6 +258,14 @@ def __enter__(self, *args):
     def __exit__(self, *args):
         self.shutdown()
 
+    @contextlib.contextmanager
+    def closing(self, *, wait=False, cancel_futures=False):
+        """Explicit context manager to do shutdown, with Wait=False by default"""
+        try:
+            yield self
+        finally:
+            self.shutdown(wait=wait, cancel_futures=cancel_futures)
+
 
 def _result_or_cancel(fut, timeout=None):
     try:

tests/test_qthreadexec.py

@@ -7,6 +7,7 @@
 import time
 import weakref
 from concurrent.futures import CancelledError, TimeoutError
+from unittest import mock
 
 import pytest
 
@@ -175,3 +176,21 @@ def func(x):
     # because the max number of workers is 5 and the rest of
     # the tasks were not started at the time of the cancel.
     assert set(results) != {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+
+
+def test_closing(executor):
+    """Test that closing context manager works as expected"""
+    # mock the shutdown method of the executor
+    with mock.patch.object(executor, "shutdown") as mock_shutdown:
+        with executor.closing():
+            pass
+
+        # ensure that shutdown was called with (False, cancel_futures=False)
+        mock_shutdown.assert_called_once_with(wait=False, cancel_futures=False)
+
+    with mock.patch.object(executor, "shutdown") as mock_shutdown:
+        with executor.closing(wait=True, cancel_futures=True):
+            pass
+
+        # ensure that shutdown was called with (False, cancel_futures=False)
+        mock_shutdown.assert_called_once_with(wait=True, cancel_futures=True)