gh-118894: Make asyncio REPL use pyrepl (GH-119433)

1 files changed, 56 insertions, 1 deletions

diff --git a/Lib/_pyrepl/console.py b/Lib/_pyrepl/console.py
index fcabf78..aa0bde8 100644
--- a/Lib/_pyrepl/console.py
+++ b/Lib/_pyrepl/console.py
@@ -19,10 +19,14 @@
 
 from __future__ import annotations
 
-import sys
+import _colorize  # type: ignore[import-not-found]
 
 from abc import ABC, abstractmethod
+import ast
+import code
 from dataclasses import dataclass, field
+import os.path
+import sys
 
 
 TYPE_CHECKING = False
@@ -136,3 +140,54 @@ class Console(ABC):
 
     @abstractmethod
     def repaint(self) -> None: ...
+
+
+class InteractiveColoredConsole(code.InteractiveConsole):
+    def __init__(
+        self,
+        locals: dict[str, object] | None = None,
+        filename: str = "<console>",
+        *,
+        local_exit: bool = False,
+    ) -> None:
+        super().__init__(locals=locals, filename=filename, local_exit=local_exit)  # type: ignore[call-arg]
+        self.can_colorize = _colorize.can_colorize()
+
+    def showsyntaxerror(self, filename=None):
+        super().showsyntaxerror(colorize=self.can_colorize)
+
+    def showtraceback(self):
+        super().showtraceback(colorize=self.can_colorize)
+
+    def runsource(self, source, filename="<input>", symbol="single"):
+        try:
+            tree = ast.parse(source)
+        except (SyntaxError, OverflowError, ValueError):
+            self.showsyntaxerror(filename)
+            return False
+        if tree.body:
+            *_, last_stmt = tree.body
+        for stmt in tree.body:
+            wrapper = ast.Interactive if stmt is last_stmt else ast.Module
+            the_symbol = symbol if stmt is last_stmt else "exec"
+            item = wrapper([stmt])
+            try:
+                code = self.compile.compiler(item, filename, the_symbol, dont_inherit=True)
+            except SyntaxError as e:
+                if e.args[0] == "'await' outside function":
+                    python = os.path.basename(sys.executable)
+                    e.add_note(
+                        f"Try the asyncio REPL ({python} -m asyncio) to use"
+                        f" top-level 'await' and run background asyncio tasks."
+                    )
+                self.showsyntaxerror(filename)
+                return False
+            except (OverflowError, ValueError):
+                self.showsyntaxerror(filename)
+                return False
+
+            if code is None:
+                return True
+
+            self.runcode(code)
+        return False