Use nursery-like concept for background tasks.

Author: jonathanslenders

prompt_toolkit/application/application.py

@@ -5,7 +5,9 @@
 import time
 from asyncio import (
     AbstractEventLoop,
+    CancelledError,
     Future,
+    Task,
     ensure_future,
     get_event_loop,
     sleep,
@@ -14,6 +16,7 @@
 from traceback import format_tb
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Dict,
     FrozenSet,
@@ -134,6 +137,10 @@ class Application(Generic[_AppResult]):
         scheduled calls), postpone the rendering max x seconds.  '0' means:
         don't postpone. '.5' means: try to draw at least twice a second.
 
+    :param refresh_interval: Automatically invalidate the UI every so many
+        seconds. When `None` (the default), only invalidate when `invalidate`
+        has been called.
+
     Filters:
 
     :param mouse_support: (:class:`~prompt_toolkit.filters.Filter` or
@@ -190,6 +197,7 @@ def __init__(self,
                  reverse_vi_search_direction: FilterOrBool = False,
                  min_redraw_interval: Union[float, int, None] = None,
                  max_render_postpone_time: Union[float, int, None] = .01,
+                 refresh_interval: Optional[float] = None,
 
                  on_reset: Optional[ApplicationEventHandler] = None,
                  on_invalidate: Optional[ApplicationEventHandler] = None,
@@ -238,6 +246,7 @@ def __init__(self,
         self.enable_page_navigation_bindings = enable_page_navigation_bindings
         self.min_redraw_interval = min_redraw_interval
         self.max_render_postpone_time = max_render_postpone_time
+        self.refresh_interval = refresh_interval
 
         # Events.
         self.on_invalidate = Event(self, on_invalidate)
@@ -386,6 +395,8 @@ def reset(self) -> None:
 
         self.exit_style = ''
 
+        self.background_tasks: List[Task] = []
+
         self.renderer.reset()
         self.key_processor.reset()
         self.layout.reset()
@@ -437,7 +448,7 @@ def schedule_redraw() -> None:
                 async def redraw_in_future() -> None:
                     await sleep(cast(float, self.min_redraw_interval) - diff)
                     schedule_redraw()
-                ensure_future(redraw_in_future())
+                self.create_background_task(redraw_in_future())
             else:
                 schedule_redraw()
         else:
@@ -488,6 +499,19 @@ def run_in_context() -> None:
         if self.context is not None:
             self.context.run(run_in_context)
 
+    def _start_auto_refresh_task(self) -> None:
+        """
+        Start a while/true loop in the background for automatic invalidation of
+        the UI.
+        """
+        async def auto_refresh():
+            while True:
+                await sleep(self.refresh_interval)
+                self.invalidate()
+
+        if self.refresh_interval:
+            self.create_background_task(auto_refresh())
+
     def _update_invalidate_events(self) -> None:
         """
         Make sure to attach 'invalidate' handlers to all invalidate events in
@@ -612,7 +636,7 @@ def read_from_input() -> None:
                     counter = flush_counter
 
                     # Automatically flush keys.
-                    ensure_future(auto_flush_input(counter))
+                    self.create_background_task(auto_flush_input(counter))
 
             async def auto_flush_input(counter: int) -> None:
                 # Flush input after timeout.
@@ -638,6 +662,7 @@ def flush_input() -> None:
                     # Draw UI.
                     self._request_absolute_cursor_position()
                     self._redraw()
+                    self._start_auto_refresh_task()
 
                     has_sigwinch = hasattr(signal, 'SIGWINCH') and in_main_thread()
                     if has_sigwinch:
@@ -696,6 +721,12 @@ async def _run_async2() -> _AppResult:
                 try:
                     result = await _run_async()
                 finally:
+                    # Wait for the background tasks to be done. This needs to
+                    # go in the finally! If `_run_async` raises
+                    # `KeyboardInterrupt`, we still want to wait for the
+                    # background tasks.
+                    await self.cancel_and_wait_for_background_tasks()
+
                     # Set the `_is_running` flag to `False`. Normally this
                     # happened already in the finally block in `run_async`
                     # above, but in case of exceptions, that's not always the
@@ -717,9 +748,8 @@ def run(self, pre_run: Optional[Callable[[], None]] = None,
         loop = get_event_loop()
 
         def run() -> _AppResult:
-            f = ensure_future(self.run_async(pre_run=pre_run))
-            get_event_loop().run_until_complete(f)
-            return f.result()
+            coro = self.run_async(pre_run=pre_run)
+            return get_event_loop().run_until_complete(coro)
 
         def handle_exception(loop, context: Dict[str, Any]) -> None:
             " Print the exception, using run_in_terminal. "
@@ -752,6 +782,32 @@ async def in_term() -> None:
         else:
             return run()
 
+    def create_background_task(self, coroutine: Awaitable[None]) -> None:
+        """
+        Start a background task (coroutine) for the running application.
+        If asyncio had nurseries like Trio, we would create a nursery in
+        `Application.run_async`, and run the given coroutine in that nursery.
+        """
+        self.background_tasks.append(get_event_loop().create_task(coroutine))
+
+    async def cancel_and_wait_for_background_tasks(self) -> None:
+        """
+        Cancel all background tasks, and wait for the cancellation to be done.
+        If any of the background tasks raised an exception, this will also
+        propagate the exception.
+
+        (If we had nurseries like Trio, this would be the `__aexit__` of a
+        nursery.)
+        """
+        for task in self.background_tasks:
+            task.cancel()
+
+        for task in self.background_tasks:
+            try:
+                await task
+            except CancelledError:
+                pass
+
     def cpr_not_supported_callback(self) -> None:
         """
         Called when we don't receive the cursor position response in time.

prompt_toolkit/buffer.py

@@ -464,7 +464,7 @@ def _text_changed(self) -> None:
         # (This happens on all change events, unlike auto completion, also when
         # deleting text.)
         if self.validator and self.validate_while_typing():
-            ensure_future(self._async_validator())
+            get_app().create_background_task(self._async_validator())
 
     def _cursor_position_changed(self) -> None:
         # Remove any complete state.
@@ -1144,11 +1144,11 @@ def insert_text(self, data: str, overwrite: bool = False,
 
             # Only complete when "complete_while_typing" is enabled.
             if self.completer and self.complete_while_typing():
-                ensure_future(self._async_completer())
+                get_app().create_background_task(self._async_completer())
 
             # Call auto_suggest.
             if self.auto_suggest:
-                ensure_future(self._async_suggester())
+                get_app().create_background_task(self._async_suggester())
 
     def undo(self) -> None:
         # Pop from the undo-stack until we find a text that if different from
@@ -1484,7 +1484,7 @@ def start_completion(
         # Only one of these options can be selected.
         assert select_first + select_last + insert_common_part <= 1
 
-        ensure_future(self._async_completer(
+        get_app().create_background_task(self._async_completer(
             select_first=select_first,
             select_last=select_last,
             insert_common_part=insert_common_part,

prompt_toolkit/contrib/ssh/server.py

@@ -58,7 +58,7 @@ def shell_requested(self) -> bool:
         return True
 
     def session_started(self) -> None:
-        asyncio.ensure_future(self._interact())
+        asyncio.get_event_loop().create_task(self._interact())
 
     async def _interact(self) -> None:
         if self._chan is None:

prompt_toolkit/history.py

@@ -9,9 +9,10 @@
 import datetime
 import os
 from abc import ABCMeta, abstractmethod
-from asyncio import ensure_future
 from typing import AsyncGenerator, Iterable, List
 
+from prompt_toolkit.application.current import get_app
+
 from .eventloop import generator_to_async_generator
 from .utils import Event
 
@@ -59,7 +60,7 @@ def start_loading(self) -> None:
         " Start loading the history. "
         if not self._loading:
             self._loading = True
-            ensure_future(self._start_loading())
+            get_app().create_background_task(self._start_loading())
 
     def get_item_loaded_event(self) -> Event['History']:
         " Event which is triggered when a new item is loaded. "

prompt_toolkit/key_binding/key_processor.py

@@ -7,7 +7,7 @@
 correct callbacks when new key presses are feed through `feed`.
 """
 import weakref
-from asyncio import ensure_future, sleep
+from asyncio import sleep
 from collections import deque
 from typing import TYPE_CHECKING, Any, Deque, Generator, List, Optional, Union
 
@@ -392,7 +392,8 @@ def _start_timeout(self) -> None:
         no key was pressed in the meantime, we flush all data in the queue and
         call the appropriate key binding handlers.
         """
-        timeout = get_app().timeoutlen
+        app = get_app()
+        timeout = app.timeoutlen
 
         if timeout is None:
             return
@@ -415,7 +416,7 @@ def flush_keys() -> None:
         # Automatically flush keys.
         # (_daemon needs to be set, otherwise, this will hang the
         # application for .5 seconds before exiting.)
-        ensure_future(wait())
+        app.create_background_task(wait())
 
 
 class KeyPressEvent:

prompt_toolkit/renderer.py

@@ -2,7 +2,7 @@
 Renders the command line on the console.
 (Redraws parts of the input line that were changed.)
 """
-from asyncio import FIRST_COMPLETED, Future, ensure_future, sleep, wait
+from asyncio import FIRST_COMPLETED, Future, sleep, wait
 from collections import deque
 from enum import Enum
 from typing import (
@@ -16,6 +16,7 @@
     Tuple,
 )
 
+from prompt_toolkit.application.current import get_app
 from prompt_toolkit.data_structures import Point, Size
 from prompt_toolkit.filters import FilterOrBool, to_filter
 from prompt_toolkit.formatted_text import AnyFormattedText, to_formatted_text
@@ -452,7 +453,7 @@ async def timer() -> None:
                             # Make sure to call this callback in the main thread.
                             self.cpr_not_supported_callback()
 
-                ensure_future(timer())
+                get_app().create_background_task(timer())
 
     def report_absolute_cursor_row(self, row: int) -> None:
         """
@@ -505,11 +506,11 @@ async def wait_for_timeout() -> None:
             # Got timeout, erase queue.
             self._waiting_for_cpr_futures = deque()
 
-        futures = [
-            ensure_future(wait_for_responses()),
-            ensure_future(wait_for_timeout()),
+        coroutines = [
+            wait_for_responses(),
+            wait_for_timeout(),
         ]
-        await wait(futures, return_when=FIRST_COMPLETED)
+        await wait(coroutines, return_when=FIRST_COMPLETED)
 
     def render(self, app: 'Application[Any]', layout: 'Layout',
                is_done: bool = False) -> None:

prompt_toolkit/shortcuts/progress_bar/base.py

@@ -7,7 +7,6 @@
         for item in pb(data):
             ...
 """
-import contextlib
 import datetime
 import functools
 import os
@@ -16,7 +15,7 @@
 import threading
 import traceback
 from asyncio import (
-    ensure_future,
+    CancelledError,
     get_event_loop,
     new_event_loop,
     set_event_loop,
@@ -190,19 +189,19 @@ def width_for_formatter(formatter: Formatter) -> AnyDimension:
             ])),
             style=self.style,
             key_bindings=self.key_bindings,
+            refresh_interval=.3,
             color_depth=self.color_depth,
             output=self.output,
             input=self.input)
 
         # Run application in different thread.
         def run() -> None:
             set_event_loop(self._app_loop)
-            with _auto_refresh_context(self.app, .3):
-                try:
-                    self.app.run()
-                except BaseException as e:
-                    traceback.print_exc()
-                    print(e)
+            try:
+                self.app.run()
+            except BaseException as e:
+                traceback.print_exc()
+                print(e)
 
         ctx: contextvars.Context = contextvars.copy_context()
 
@@ -230,6 +229,7 @@ def __exit__(self, *a: object) -> None:
 
         if self._thread is not None:
             self._thread.join()
+        self._app_loop.close()
 
     def __call__(self,
                  data: Optional[Iterable[_T]] = None,
@@ -372,28 +372,3 @@ def time_left(self) -> Optional[datetime.timedelta]:
             return None
         else:
             return self.time_elapsed * (100 - self.percentage) / self.percentage
-
-
-@contextlib.contextmanager
-def _auto_refresh_context(
-        app: 'Application', refresh_interval: Optional[float] = None
-) -> Generator[None, None, None]:
-    """
-    Return a context manager for the auto-refresh loop.
-    """
-    done = False
-
-    async def run() -> None:
-        if refresh_interval:
-            while not done:
-                await sleep(refresh_interval)
-                app.invalidate()
-
-    if refresh_interval:
-        ensure_future(run())
-
-    try:
-        yield
-    finally:
-        # Exit.
-        done = True