diff --git a/fire/core.py b/fire/core.py index 8e23e76b..81755393 100644 --- a/fire/core.py +++ b/fire/core.py @@ -50,6 +50,7 @@ def main(argv): """ import asyncio +import concurrent.futures import inspect import json import os @@ -649,6 +650,30 @@ def _GetMember(component, args): raise FireError('Could not consume arg:', arg) +def _RunCoroutine(coro): + """Runs a coroutine, blocking until it completes. + + Args: + coro: The coroutine to run. + + Returns: + The value the coroutine returns. + + Raises: + Exception: Whatever the coroutine raises. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro) + # asyncio.run refuses to be called while this thread already runs a loop, + # which happens when Fire is called from async code. Drive the coroutine + # from a loop belonging to a worker thread instead, so the call still + # blocks here until the coroutine completes. + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, coro).result() + + def _CallAndUpdateTrace(component, args, component_trace, treatment='class', target=None): """Call the component by consuming args from args, and update the FireTrace. @@ -678,14 +703,7 @@ def _CallAndUpdateTrace(component, args, component_trace, treatment='class', # Call the function. if inspectutils.IsCoroutineFunction(fn): - try: - loop = asyncio.get_running_loop() - except RuntimeError: - # No event loop running, create a new one - component = asyncio.run(fn(*varargs, **kwargs)) - else: - # Event loop is already running - component = loop.run_until_complete(fn(*varargs, **kwargs)) + component = _RunCoroutine(fn(*varargs, **kwargs)) else: component = fn(*varargs, **kwargs) diff --git a/fire/fire_test.py b/fire/fire_test.py index 99b4a7c6..6956a72c 100644 --- a/fire/fire_test.py +++ b/fire/fire_test.py @@ -14,6 +14,7 @@ """Tests for the fire module.""" +import asyncio import os import sys from unittest import mock @@ -716,6 +717,23 @@ def testFireAsyncio(self): self.assertEqual(fire.Fire(tc.py3.WithAsyncio, command=['double', '--count', '10']), 20) + def testFireAsyncioFromRunningLoop(self): + async def main(): + return fire.Fire(tc.py3.WithAsyncio, + command=['double', '--count', '10']) + + self.assertEqual(asyncio.run(main()), 20) + + def testFireAsyncioFromRunningLoopRaises(self): + async def raiser(count=0): + raise ValueError('boom') + + async def main(): + return fire.Fire(raiser, command=['--count', '3']) + + with self.assertRaises(ValueError): + asyncio.run(main()) + if __name__ == '__main__': testutils.main()