Skip to content

Commit

Permalink
Many asyncio REPL improvements.
Browse files Browse the repository at this point in the history
- Added `--asyncio` flag to the `ptpython` entry point to activate the
  asyncio-REPL. This will ensure that an event loop is created at the start in
  which we can run top-level await statements.
- Use `get_running_loop()` instead of `get_event_loop()`.
- Better handling of `SystemExit` and control-c in the async REPL.
  • Loading branch information
jonathanslenders committed Dec 12, 2023
1 parent 6801f94 commit f019301
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 25 deletions.
2 changes: 1 addition & 1 deletion ptpython/contrib/asyncssh_repl.py
Expand Up @@ -110,7 +110,7 @@ def data_received(self, data: AnyStr, datatype: int | None) -> None:
"""
When data is received, send to inputstream of the CLI and repaint.
"""
self._input_pipe.send(data)
self._input_pipe.send(data) # type: ignore

def _print(
self, *data: object, sep: str = " ", end: str = "\n", file: Any = None
Expand Down
13 changes: 12 additions & 1 deletion ptpython/entry_points/run_ptpython.py
Expand Up @@ -9,6 +9,7 @@
-h, --help show this help message and exit
--vi Enable Vi key bindings
-i, --interactive Start interactive shell after executing this file.
--asyncio Run an asyncio event loop to support top-level "await".
--light-bg Run on a light background (use dark colors for text).
--dark-bg Run on a dark background (use light colors for text).
--config-file CONFIG_FILE
Expand All @@ -24,6 +25,7 @@
from __future__ import annotations

import argparse
import asyncio
import os
import pathlib
import sys
Expand Down Expand Up @@ -68,6 +70,11 @@ def create_parser() -> _Parser:
action="store_true",
help="Start interactive shell after executing this file.",
)
parser.add_argument(
"--asyncio",
action="store_true",
help='Run an asyncio event loop to support top-level "await".',
)
parser.add_argument(
"--light-bg",
action="store_true",
Expand Down Expand Up @@ -206,16 +213,20 @@ def configure(repl: PythonRepl) -> None:

import __main__

embed(
embed_result = embed( # type: ignore
vi_mode=a.vi,
history_filename=history_file,
configure=configure,
locals=__main__.__dict__,
globals=__main__.__dict__,
startup_paths=startup_paths,
title="Python REPL (ptpython)",
return_asyncio_coroutine=a.asyncio,
)

if a.asyncio:
asyncio.run(embed_result)


if __name__ == "__main__":
run()
4 changes: 2 additions & 2 deletions ptpython/python_input.py
Expand Up @@ -4,7 +4,7 @@
"""
from __future__ import annotations

from asyncio import get_event_loop
from asyncio import get_running_loop
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Mapping, TypeVar, Union

Expand Down Expand Up @@ -1010,7 +1010,7 @@ def get_signatures_in_executor(document: Document) -> list[Signature]:
app = self.app

async def on_timeout_task() -> None:
loop = get_event_loop()
loop = get_running_loop()

# Never run multiple get-signature threads.
if self._get_signatures_thread_running:
Expand Down
76 changes: 55 additions & 21 deletions ptpython/repl.py
Expand Up @@ -12,6 +12,7 @@
import asyncio
import builtins
import os
import signal
import sys
import traceback
import types
Expand Down Expand Up @@ -158,27 +159,58 @@ def run(self) -> None:
clear_title()
self._remove_from_namespace()

async def run_and_show_expression_async(self, text: str) -> object:
loop = asyncio.get_event_loop()
async def run_and_show_expression_async(self, text: str) -> Any:
loop = asyncio.get_running_loop()
system_exit: SystemExit | None = None

try:
result = await self.eval_async(text)
except KeyboardInterrupt: # KeyboardInterrupt doesn't inherit from Exception.
raise
except SystemExit:
return
except BaseException as e:
self._handle_exception(e)
else:
# Print.
if result is not None:
await loop.run_in_executor(None, lambda: self._show_result(result))
try:
# Create `eval` task. Ensure that control-c will cancel this
# task.
async def eval() -> Any:
nonlocal system_exit
try:
return await self.eval_async(text)
except SystemExit as e:
# Don't propagate SystemExit in `create_task()`. That
# will kill the event loop. We want to handle it
# gracefully.
system_exit = e

task = asyncio.create_task(eval())
loop.add_signal_handler(signal.SIGINT, lambda *_: task.cancel())
result = await task

if system_exit is not None:
raise system_exit
except KeyboardInterrupt:
# KeyboardInterrupt doesn't inherit from Exception.
raise
except SystemExit:
raise
except BaseException as e:
self._handle_exception(e)
else:
# Print.
if result is not None:
await loop.run_in_executor(None, lambda: self._show_result(result))

# Loop.
self.current_statement_index += 1
self.signatures = []
# Return the result for future consumers.
return result
# Loop.
self.current_statement_index += 1
self.signatures = []
# Return the result for future consumers.
return result
finally:
loop.remove_signal_handler(signal.SIGINT)

except KeyboardInterrupt as e:
# Handle all possible `KeyboardInterrupt` errors. This can
# happen during the `eval`, but also during the
# `show_result` if something takes too long.
# (Try/catch is around the whole block, because we want to
# prevent that a Control-C keypress terminates the REPL in
# any case.)
self._handle_keyboard_interrupt(e)

async def run_async(self) -> None:
"""
Expand All @@ -192,7 +224,7 @@ async def run_async(self) -> None:
(Both for control-C to work, as well as for the code to see the right
thread in which it was embedded).
"""
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()

if self.terminal_title:
set_title(self.terminal_title)
Expand Down Expand Up @@ -222,6 +254,8 @@ async def run_async(self) -> None:
# `KeyboardInterrupt` exceptions can end up in the event
# loop selector.
self._handle_keyboard_interrupt(e)
except SystemExit:
return
finally:
if self.terminal_title:
clear_title()
Expand Down Expand Up @@ -250,7 +284,7 @@ def eval(self, line: str) -> object:
result = eval(code, self.get_globals(), self.get_locals())

if _has_coroutine_flag(code):
result = asyncio.get_event_loop().run_until_complete(result)
result = asyncio.get_running_loop().run_until_complete(result)

self._store_eval_result(result)
return result
Expand All @@ -263,7 +297,7 @@ def eval(self, line: str) -> object:
result = eval(code, self.get_globals(), self.get_locals())

if _has_coroutine_flag(code):
result = asyncio.get_event_loop().run_until_complete(result)
result = asyncio.get_running_loop().run_until_complete(result)

return None

Expand Down

0 comments on commit f019301

Please sign in to comment.