Improved usage of Application background tasks.

Author: jonathanslenders

prompt_toolkit/application/application.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 import re
 import signal
@@ -782,13 +783,15 @@ async def in_term() -> None:
  else:
  return run()
 
- def create_background_task(self, coroutine: Awaitable[None]) -> None:
+ def create_background_task(self, coroutine: Awaitable[None]) -> 'asyncio.Task[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))
+ task = get_event_loop().create_task(coroutine)
+ self.background_tasks.append(task)
+ return task
 
  async def cancel_and_wait_for_background_tasks(self) -> None:
  """

prompt_toolkit/buffer.py

@@ -2,12 +2,13 @@
 Data structures for the Buffer.
 It holds the text, cursor position, history, etc...
 """
+import asyncio
 import os
 import re
 import shlex
 import subprocess
 import tempfile
-from asyncio import Future, ensure_future
+from asyncio import Task, ensure_future
 from enum import Enum
 from functools import wraps
 from typing import (
@@ -1387,7 +1388,7 @@ def apply_search(self, search_state: SearchState,
  def exit_selection(self) -> None:
  self.selection_state = None
 
- def open_in_editor(self, validate_and_handle: bool = False) -> 'Future[None]':
+ def open_in_editor(self, validate_and_handle: bool = False) -> 'asyncio.Task[None]':
  """
  Open code in editor.
 
@@ -1432,7 +1433,7 @@ async def run() -> None:
  # Clean up temp file.
  os.remove(filename)
 
- return ensure_future(run())
+ return get_app().create_background_task(run())
 
  def _open_file_in_editor(self, filename: str) -> bool:
  """

prompt_toolkit/contrib/telnet/server.py

@@ -1,9 +1,10 @@
 """
 Telnet server.
 """
+import asyncio
 import contextvars # Requires Python3.7!
 import socket
-from asyncio import ensure_future, get_event_loop
+from asyncio import get_event_loop
 from typing import (
  Awaitable,
  Callable,
@@ -244,6 +245,7 @@ def __init__(self, host: str = '127.0.0.1', port: int = 23,
  self.interact = interact
  self.encoding = encoding
  self.style = style
+ self._application_tasks: List[asyncio.Task] = []
 
  self.connections: Set[TelnetConnection] = set()
  self._listen_socket: Optional[socket.socket] = None
@@ -268,11 +270,18 @@ def start(self) -> None:
 
  get_event_loop().add_reader(self._listen_socket, self._accept)
 
- def stop(self) -> None:
+ async def stop(self) -> None:
  if self._listen_socket:
  get_event_loop().remove_reader(self._listen_socket)
  self._listen_socket.close()
 
+ # Wait for all applications to finish.
+ for t in self._application_tasks:
+ t.cancel()
+
+ for t in self._application_tasks:
+ await t
+
  def _accept(self) -> None:
  """
  Accept new incoming connection.
@@ -297,6 +306,8 @@ async def run() -> None:
  print(e)
  finally:
  self.connections.remove(connection)
+ self._application_tasks.remove(task)
  logger.info('Stopping interaction %r %r', *addr)
 
- ensure_future(run())
+ task = get_event_loop().create_task(run())
+ self._application_tasks.append(task)

prompt_toolkit/key_binding/bindings/completion.py

@@ -1,8 +1,8 @@
 """
 Key binding handlers for displaying completions.
 """
+import asyncio
 import math
-from asyncio import Future, ensure_future
 from typing import TYPE_CHECKING, List
 
 from prompt_toolkit.application.run_in_terminal import in_terminal
@@ -80,7 +80,7 @@ def display_completions_like_readline(event: E) -> None:
 
 
 def _display_completions_like_readline(
- app: 'Application', completions: List[Completion]) -> 'Future[None]':
+ app: 'Application', completions: List[Completion]) -> 'asyncio.Task[None]':
  """
  Display the list of completions in columns above the prompt.
  This will ask for a confirmation if there are too many completions to fit
@@ -159,7 +159,7 @@ async def run_compl() -> None:
  # Display all completions.
  display(0)
 
- return ensure_future(run_compl())
+ return app.create_background_task(run_compl())
 
 
 def _create_more_session(message: str = '--MORE--') -> 'PromptSession':

prompt_toolkit/shortcuts/prompt.py

@@ -24,7 +24,7 @@
  s = PromptSession()
  result = s.prompt('Say something: ')
 """
-from asyncio import ensure_future, get_event_loop, sleep
+from asyncio import get_event_loop, sleep
 from enum import Enum
 from functools import partial
 from typing import (