8000 Improved usage of Application background tasks. · cool-RR/python-prompt-toolkit@ba49520 · GitHub
[go: up one dir, main page]

Skip to content

Commit ba49520

Browse files
Improved usage of Application background tasks.
1 parent 4023ad9 commit ba49520

File tree

5 files changed

+27
-12
lines changed

5 files changed

+27
-12
lines changed

prompt_toolkit/application/application.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
import os
23
import re
34
import signal
@@ -782,13 +783,15 @@ async def in_term() -> None:
782783
else:
783784
return run()
784785

785-
def create_background_task(self, coroutine: Awaitable[None]) -> None:
786+
def create_background_task(self, coroutine: Awaitable[None]) -> 'asyncio.Task[None]':
786787
"""
787788
Start a background task (coroutine) for the running application.
788789
If asyncio had nurseries like Trio, we would create a nursery in
789790
`Application.run_async`, and run the given coroutine in that nursery.
790791
"""
791-
self.background_tasks.append(get_event_loop().create_task(coroutine))
792+
task = get_event_loop().create_task(coroutine)
793+
self.background_tasks.append(task)
794+
return task
792795

793796
async def cancel_and_wait_for_background_tasks(self) -> None:
794797
"""

prompt_toolkit/buffer.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,13 @@
22
Data structures for the Buffer.
33
It holds the text, cursor position, history, etc...
44
"""
5+
import asyncio
56
import os
67
import re
78
import shlex
89
import subprocess
910
import tempfile
10-
from asyncio import Future, ensure_future
11+
from asyncio import Task, ensure_future
1112
from enum import Enum
1213
from functools import wraps
1314
from typing import (
@@ -1387,7 +1388,7 @@ def apply_search(self, search_state: SearchState,
13871388
def exit_selection(self) -> None:
13881389
self.selection_state = None
13891390

1390-
def open_in_editor(self, validate_and_handle: bool = False) -> 'Future[None]':
1391+
def open_in_editor(self, validate_and_handle: bool = False) -> 'asyncio.Task[None]':
13911392
"""
13921393
Open code in editor.
13931394
@@ -1432,7 +1433,7 @@ async def run() -> None:
14321433
# Clean up temp file.
14331434
os.remove(filename)
14341435

1435-
return ensure_future(run())
1436+
return get_app().create_background_task(run())
14361437

14371438
def _open_file_in_editor(self, filename: str) -> bool:
14381439
"""

prompt_toolkit/contrib/telnet/server.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
"""
22
Telnet server.
33
"""
4+
import asyncio
45
import contextvars # Requires Python3.7!
56
import socket
6-
from asyncio import ensure_future, get_event_loop
7+
from asyncio import get_event_loop
78
from typing import (
89
Awaitable,
910
Callable,
@@ -244,6 +245,7 @@ def __init__(self, host: str = '127.0.0.1', port: int = 23,
244245
self.interact = interact
245246
self.encoding = encoding
246247
self.style = style
248+
self._application_tasks: List[asyncio.Task] = []
247249

248250
self.connections: Set[TelnetConnection] = set()
249251
self._listen_socket: Optional[socket.socket] = None
@@ -268,11 +270,18 @@ def start(self) -> None:
268270

269271
get_event_loop().add_reader(self._listen_socket, self._accept)
270272

271-
def stop(self) -> None:
273+
async def stop(self) -> None:
272274
if self._listen_socket:
273275
get_event_loop().remove_reader(self._listen_socket)
274276
self._listen_socket.close()
275277

278+
# Wait for all applications to finish.
279+
for t in self._application_tasks:
280+
t.cancel()
281+
282+
for t in self._application_tasks:
283+
await t
284+
276285
def _accept(self) -> None:
277286
"""
278287
Accept new incoming connection.
@@ -297,6 +306,8 @@ async def run() -> None:
297306
print(e)
298307
finally:
299308
self.connections.remove(connection)
309+
self._application_tasks.remove(task)
300310
logger.info('Stopping interaction %r %r', *addr)
301311

302-
ensure_future(run())
312+
task = get_event_loop().create_task(run())
313+
self._application_tasks.append(task)

prompt_toolkit/key_binding/bindings/completion.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
"""
22
Key binding handlers for displaying completions.
33
"""
4+
import asyncio
45
import math
5-
from asyncio import Future, ensure_future
66
from typing import TYPE_CHECKING, List
77

88
from prompt_toolkit.application.run_in_terminal import in_terminal
@@ -80,7 +80,7 @@ def display_completions_like_readline(event: E) -> None:
8080

8181

8282
def _display_completions_like_readline(
83-
app: 'Application', completions: List[Completion]) -> 'Future[None]':
83+
app: 'Application', completions: List[Completion]) -> 'asyncio.Task[None]':
8484
"""
8585
Display the list of completions in columns above the prompt.
8686
This will ask for a confirmation if there are too many completions to fit
@@ -159,7 +159,7 @@ async def run_compl() -> None:
159159
# Display all completions.
160160
display(0)
161161

162-
return ensure_future(run_compl())
162+
return app.create_background_task(run_compl())
163163

164164

165165
def _create_more_session(message: str = '--MORE--') -> 'PromptSession':

prompt_toolkit/shortcuts/prompt.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
s = PromptSession()
2525
result = s.prompt('Say something: ')
2626
"""
27-
from asyncio import ensure_future, get_event_loop, sleep
27+
from asyncio import get_event_loop, sleep
2828
from enum import Enum
2929
from functools import partial
3030
from typing import (

0 commit comments

Comments
 (0)
0