From d1a1bca1f0550a4715f1bf32b1586caa7bc4487b Mon Sep 17 00:00:00 2001
From: Dino Viehland <dinoviehland@meta.com>
Date: Tue, 30 Jul 2024 05:03:52 -0700
Subject: [PATCH] gh-119896: Fix CTRL-Z behavior in the new REPL on Windows
 (GH-122217)

---
 Lib/_pyrepl/reader.py          | 9 +++++++--
 Lib/_pyrepl/simple_interact.py | 1 +
 Lib/_pyrepl/utils.py           | 3 ++-
 Lib/_pyrepl/windows_console.py | 5 ++++-
 4 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/Lib/_pyrepl/reader.py b/Lib/_pyrepl/reader.py
index 8b282a382d3..13b1f3eb9d1 100644
--- a/Lib/_pyrepl/reader.py
+++ b/Lib/_pyrepl/reader.py
@@ -21,6 +21,8 @@
 
 from __future__ import annotations
 
+import sys
+
 from contextlib import contextmanager
 from dataclasses import dataclass, field, fields
 import unicodedata
@@ -52,7 +54,10 @@ def disp_str(buffer: str) -> tuple[str, list[int]]:
     b: list[int] = []
     s: list[str] = []
     for c in buffer:
-        if ord(c) < 128:
+        if c == '\x1a':
+            s.append(c)
+            b.append(2)
+        elif ord(c) < 128:
             s.append(c)
             b.append(1)
         elif unicodedata.category(c).startswith("C"):
@@ -110,7 +115,7 @@ default_keymap: tuple[tuple[KeySpec, CommandName], ...] = tuple(
         (r"\C-w", "unix-word-rubout"),
         (r"\C-x\C-u", "upcase-region"),
         (r"\C-y", "yank"),
-        (r"\C-z", "suspend"),
+        *(() if sys.platform == "win32" else ((r"\C-z", "suspend"), )),
         (r"\M-b", "backward-word"),
         (r"\M-c", "capitalize-word"),
         (r"\M-d", "kill-word"),
diff --git a/Lib/_pyrepl/simple_interact.py b/Lib/_pyrepl/simple_interact.py
index 2c3dffe070c..91aef5e01eb 100644
--- a/Lib/_pyrepl/simple_interact.py
+++ b/Lib/_pyrepl/simple_interact.py
@@ -76,6 +76,7 @@ REPL_COMMANDS = {
     "copyright": _sitebuiltins._Printer('copyright', sys.copyright),
     "help": "help",
     "clear": _clear_screen,
+    "\x1a": _sitebuiltins.Quitter('\x1a', ''),
 }
 
 
diff --git a/Lib/_pyrepl/utils.py b/Lib/_pyrepl/utils.py
index 20dbb1f7e17..0f36083b6ff 100644
--- a/Lib/_pyrepl/utils.py
+++ b/Lib/_pyrepl/utils.py
@@ -21,4 +21,5 @@ def wlen(s: str) -> int:
     length = sum(str_width(i) for i in s)
     # remove lengths of any escape sequences
     sequence = ANSI_ESCAPE_SEQUENCE.findall(s)
-    return length - sum(len(i) for i in sequence)
+    ctrl_z_cnt = s.count('\x1a')
+    return length - sum(len(i) for i in sequence) + ctrl_z_cnt
diff --git a/Lib/_pyrepl/windows_console.py b/Lib/_pyrepl/windows_console.py
index 9e97b1524e2..ba9af36b8be 100644
--- a/Lib/_pyrepl/windows_console.py
+++ b/Lib/_pyrepl/windows_console.py
@@ -253,7 +253,7 @@ class WindowsConsole(Console):
         else:
             self.__posxy = wlen(newline), y
 
-            if "\x1b" in newline or y != self.__posxy[1]:
+            if "\x1b" in newline or y != self.__posxy[1] or '\x1a' in newline:
                 # ANSI escape characters are present, so we can't assume
                 # anything about the position of the cursor.  Moving the cursor
                 # to the left margin should work to get to a known position.
@@ -291,6 +291,9 @@ class WindowsConsole(Console):
         self.__write("\x1b[?12l")
 
     def __write(self, text: str) -> None:
+        if "\x1a" in text:
+            text = ''.join(["^Z" if x == '\x1a' else x for x in text])
+
         if self.out is not None:
             self.out.write(text.encode(self.encoding, "replace"))
             self.out.flush()