backport Runner for test_asyncio.test_subprocess

Author: graingert

Lib/test/test_asyncio/test_subprocess.py

@@ -1,12 +1,18 @@
+importcontextvars
+importenum
+importfunctools
 importos
 importsignal
 importsys
+importthreading
 importunittest
 importwarnings
 fromunittestimportmock
 
-importasyncio
+importasyncio.runners
 fromasyncioimportbase_subprocess
+fromasyncioimportcoroutines
+fromasyncioimportevents
 fromasyncioimportsubprocess
 fromtest.test_asyncioimportutilsastest_utils
 fromtestimportsupport
@@ -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"
+
+classRunner:
+"""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()
+returnself
+
+def__exit__(self, exc_type, exc_val, exc_tb):
+self.close()
+
+defclose(self):
+"""Shutdown and close event loop."""
+ifself._stateisnot_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:
+ifself._set_event_loop:
+events.set_event_loop(None)
+loop.close()
+self._loop=None
+self._state=_State.CLOSED
+
+defget_loop(self):
+"""Return embedded event loop."""
+self._lazy_init()
+returnself._loop
+
+defrun(self, coro, *, context=None):
+"""Run a coroutine inside the embedded event loop."""
+ifnotcoroutines.iscoroutine(coro):
+raiseValueError("a coroutine was expected, got {!r}".format(coro))
+
+ifevents._get_running_loop() isnotNone:
+# fail fast with short traceback
+raiseRuntimeError(
+"Runner.run() cannot be called from a running event loop")
+
+self._lazy_init()
+
+ifcontextisNone:
+context=self._context
+task=self._loop.create_task(coro)
+
+if (threading.current_thread() isthreading.main_thread()
+andsignal.getsignal(signal.SIGINT) issignal.default_int_handler
+ ):
+sigint_handler=functools.partial(self._on_sigint, main_task=task)
+try:
+signal.signal(signal.SIGINT, sigint_handler)
+exceptValueError:
+# `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:
+ifself._set_event_loop:
+events.set_event_loop(self._loop)
+returnself._loop.run_until_complete(task)
+exceptexceptions.CancelledError:
+ifself._interrupt_count>0andtask.uncancel() ==0:
+raiseKeyboardInterrupt()
+else:
+raise# CancelledError
+finally:
+if (sigint_handlerisnotNone
+andsignal.getsignal(signal.SIGINT) issigint_handler
+ ):
+signal.signal(signal.SIGINT, signal.default_int_handler)
+
+def_lazy_init(self):
+ifself._stateis_State.CLOSED:
+raiseRuntimeError("Runner is closed")
+ifself._stateis_State.INITIALIZED:
+return
+ifself._loop_factoryisNone:
+self._loop=events.new_event_loop()
+self._set_event_loop=True
+else:
+self._loop=self._loop_factory()
+ifself._debugisnotNone:
+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
+ifself._interrupt_count==1andnotmain_task.done():
+main_task.cancel()
+# wakeup loop if it is blocked by select() with long timeout
+self._loop.call_soon_threadsafe(lambda: None)
+return
+raiseKeyboardInterrupt()
+
+
 ifsys.platform!='win32':
 # Unix
 classSubprocessWatcherMixin(SubprocessMixin):
@@ -717,7 +864,7 @@ async def execute():
 
 watcher.add_child_handler.assert_not_called()
 
-withasyncio.Runner(loop_factory=asyncio.new_event_loop) asrunner:
+withRunner(loop_factory=asyncio.new_event_loop) asrunner:
 self.assertIsNone(runner.run(execute()))
 self.assertListEqual(watcher.mock_calls, [
 mock.call.__enter__(),