backport Runner for test_asyncio.test_subprocess

graingert · Jul 21, 2022 · Jul 21, 2022 · Jul 21, 2022 · ca487621b546fbd2a1fa3425c74c35fb76cc84c1
commit ca487621b546fbd2a1fa3425c74c35fb76cc84c1
@@ -1,12 +1,18 @@
+import contextvars
+import enum
+import functools
 import os
 import signal
 import sys
+import threading
 import unittest
 import warnings
 from unittest import mock
 
-import asyncio
+import asyncio.runners
 from asyncio import base_subprocess
+from asyncio import coroutines
+from asyncio import events
 from asyncio import subprocess
 from test.test_asyncio import utils as test_utils
 from test import support
@@ -636,6 +642,147 @@ async def execute():
         self.assertIsNone(self.loop.run_until_complete(execute()))
 
 
+class _State(enum.Enum):
+    CREATED = "created"
+    INITIALIZED = "initialized"
+    CLOSED = "closed"
+
+class Runner:
+    """A context manager that controls event loop life cycle.
+
+    The context manager always creates a new event loop,
+    allows to run async functions inside it,
+    and properly finalizes the loop at the context manager exit.
+
+    If debug is True, the event loop will be run in debug mode.
+    If loop_factory is passed, it is used for new event loop creation.
+
+    asyncio.run(main(), debug=True)
+
+    is a shortcut for
+
+    with asyncio.Runner(debug=True) as runner:
+        runner.run(main())
+
+    The run() method can be called multiple times within the runner's context.
+
+    This can be useful for interactive console (e.g. IPython),
+    unittest runners, console tools, -- everywhere when async code
+    is called from existing sync framework and where the preferred single
+    asyncio.run() call doesn't work.
+
+    """
+
+    # Note: the class is final, it is not intended for inheritance.
+
+    def __init__(self, *, debug=None, loop_factory=None):
+        self._state = _State.CREATED
+        self._debug = debug
+        self._loop_factory = loop_factory
+        self._loop = None
+        self._context = None
+        self._interrupt_count = 0
+        self._set_event_loop = False
+
+    def __enter__(self):
+        self._lazy_init()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def close(self):
+        """Shutdown and close event loop."""
+        if self._state is not _State.INITIALIZED:
+            return
+        try:
+            loop = self._loop
+            asyncio.runners._cancel_all_tasks(loop)
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.run_until_complete(loop.shutdown_default_executor())
+        finally:
+            if self._set_event_loop:
+                events.set_event_loop(None)
+            loop.close()
+            self._loop = None
+            self._state = _State.CLOSED
+
+    def get_loop(self):
+        """Return embedded event loop."""
+        self._lazy_init()
+        return self._loop
+
+    def run(self, coro, *, context=None):
+        """Run a coroutine inside the embedded event loop."""
+        if not coroutines.iscoroutine(coro):
+            raise ValueError("a coroutine was expected, got {!r}".format(coro))
+
+        if events._get_running_loop() is not None:
+            # fail fast with short traceback
+            raise RuntimeError(
+                "Runner.run() cannot be called from a running event loop")
+
+        self._lazy_init()
+
+        if context is None:
+            context = self._context
+        task = self._loop.create_task(coro)
+
+        if (threading.current_thread() is threading.main_thread()
+            and signal.getsignal(signal.SIGINT) is signal.default_int_handler
+        ):
+            sigint_handler = functools.partial(self._on_sigint, main_task=task)
+            try:
+                signal.signal(signal.SIGINT, sigint_handler)
+            except ValueError:
+                # `signal.signal` may throw if `threading.main_thread` does
+                # not support signals (e.g. embedded interpreter with signals
+                # not registered - see gh-91880)
+                sigint_handler = None
+        else:
+            sigint_handler = None
+
+        self._interrupt_count = 0
+        try:
+            if self._set_event_loop:
+                events.set_event_loop(self._loop)
+            return self._loop.run_until_complete(task)
+        except exceptions.CancelledError:
+            if self._interrupt_count > 0 and task.uncancel() == 0:
+                raise KeyboardInterrupt()
+            else:
+                raise  # CancelledError
+        finally:
+            if (sigint_handler is not None
+                and signal.getsignal(signal.SIGINT) is sigint_handler
+            ):
+                signal.signal(signal.SIGINT, signal.default_int_handler)
+
+    def _lazy_init(self):
+        if self._state is _State.CLOSED:
+            raise RuntimeError("Runner is closed")
+        if self._state is _State.INITIALIZED:
+            return
+        if self._loop_factory is None:
+            self._loop = events.new_event_loop()
+            self._set_event_loop = True
+        else:
+            self._loop = self._loop_factory()
+        if self._debug is not None:
+            self._loop.set_debug(self._debug)
+        self._context = contextvars.copy_context()
+        self._state = _State.INITIALIZED
+
+    def _on_sigint(self, signum, frame, main_task):
+        self._interrupt_count += 1
+        if self._interrupt_count == 1 and not main_task.done():
+            main_task.cancel()
+            # wakeup loop if it is blocked by select() with long timeout
+            self._loop.call_soon_threadsafe(lambda: None)
+            return
+        raise KeyboardInterrupt()
+
+
 if sys.platform != 'win32':
     # Unix
     class SubprocessWatcherMixin(SubprocessMixin):
@@ -717,7 +864,7 @@ async def execute():
 
                 watcher.add_child_handler.assert_not_called()
 
-            with asyncio.Runner(loop_factory=asyncio.new_event_loop) as runner:
+            with Runner(loop_factory=asyncio.new_event_loop) as runner:
                 self.assertIsNone(runner.run(execute()))
             self.assertListEqual(watcher.mock_calls, [
                 mock.call.__enter__(),