Add asyncio.forever()

python · 1st1 · Nov 10, 2016 · Nov 15, 2016 · Nov 15, 2016 · Nov 15, 2016
commit 9acdcebe7f81acbecd86ed1d2247cdb4628a7455
diff --git a/asyncio/__init__.py b/asyncio/__init__.py
@@ -41,7 +41,8 @@
            streams.__all__ +
            subprocess.__all__ +
            tasks.__all__ +
-           transports.__all__)
+           transports.__all__ +
+           ['run', 'forever'])  # Will fix this later.
 
 if sys.platform == 'win32':  # pragma: no cover
     from .windows_events import *

diff --git a/asyncio/base_events.py b/asyncio/base_events.py
@@ -294,6 +294,9 @@ def __init__(self):
         # Set to True when `loop.shutdown_asyncgens` is called.
         self._asyncgens_shutdown_called = False
 
+        # Future that isn't resolved while the loop is running.
+        self._forever_fut = None
+
     def __repr__(self):
         return ('<%s running=%s closed=%s debug=%s>'
                 % (self.__class__.__name__, self.is_running(),
@@ -430,8 +433,12 @@ def shutdown_asyncgens(self):
                     'asyncgen': agen
                 })
 
+    def get_forever_future(self):
+        return self._forever_fut
+
     def run_forever(self):
         """Run until stop() is called."""
+        self._forever_fut = self.create_future()
         self._check_closed()
         if self.is_running():
             raise RuntimeError('This event loop is already running')
@@ -450,7 +457,14 @@ def run_forever(self):
                 self._run_once()
                 if self._stopping:
                     break
+        except BaseException as ex:
+            self._forever_fut.set_exception(ex)
+            self._forever_fut._log_traceback = False
+            raise ex
+        else:
+            self._forever_fut.set_result(None)
         finally:
+            self._forever_fut = None
             self._stopping = False
             self._thread_id = None
             events._set_running_loop(None)

diff --git a/asyncio/events.py b/asyncio/events.py
@@ -512,6 +512,9 @@ def get_debug(self):
     def set_debug(self, enabled):
         raise NotImplementedError
 
+    def get_forever_future(self):
+        raise NotImplementedError
+
 
 class AbstractEventLoopPolicy:
     """Abstract policy for accessing the event loop."""

diff --git a/asyncio/run.py b/asyncio/run.py
@@ -1,13 +1,40 @@
 """asyncio.run() function."""
 
-__all__ = ['run']
+__all__ = ['run', 'forever']
 
 import threading
 
 from . import coroutines
 from . import events
 
 
+@coroutines.coroutine
+def forever():
+    """Wait until the current event loop stops running.
+
+    The coroutine will return None if the loop is stopped by
+    calling the `loop.stop()` method.
+
+    The coroutine will propagate any exception that caused
+    the loop to stop;
+
+    It is recommended to use this coroutine with the asyncio.run()
+    function:
+
+        async def coro():
+            print('hi')
+            try:
+                await asyncio.forever()
+            except KeyboardInterrupt:
+                await asyncio.sleep(1)
+                print('bye')
+
+        asyncio.run(coro())
+    """
+    loop = events.get_event_loop()
+    return (yield from loop.get_forever_future())
+
+
 def run(coro, *, debug=False):
     """Run a coroutine.
 
@@ -50,7 +77,15 @@ async def main():
         if debug:
             loop.set_debug(True)
 
-        result = loop.run_until_complete(coro)
+        task = loop.create_task(coro)
+        task.add_done_callback(lambda task: loop.stop())
+
+        try:
+            loop.run_forever()
+        except BaseException as ex:
+            result = loop.run_until_complete(task)
+        else:
+            result = task.result()
 
         try:
             # `shutdown_asyncgens` was added in Python 3.6; not all

diff --git a/asyncio/tasks.py b/asyncio/tasks.py
@@ -302,7 +302,7 @@ def _step(self, exc=None):
     def _wakeup(self, future):
         try:
             future.result()
-        except Exception as exc:
+        except BaseException as exc:
             # This may also be a cancellation.
             self._step(exc)
         else: