Many asyncio REPL improvements.

- 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.

Author: jonathanslenders

ptpython/contrib/asyncssh_repl.py

@@ -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

ptpython/entry_points/run_ptpython.py

@@ -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
@@ -24,6 +25,7 @@
 from __future__ import annotations
 
 import argparse
+import asyncio
 import os
 import pathlib
 import sys
@@ -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",
@@ -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()

ptpython/python_input.py

@@ -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
 
@@ -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:

ptpython/repl.py

@@ -12,6 +12,7 @@
 import asyncio
 import builtins
 import os
+import signal
 import sys
 import traceback
 import types
@@ -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:
         """
@@ -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)
@@ -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()
@@ -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
@@ -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