gh-150582: Use _winapi for _pyrepl console mode and input flush

Author: aisk

Lib/_pyrepl/windows_console.py

@@ -165,19 +165,16 @@ def __init__(
             trace('console supports virtual terminal')
 
         # Save original console modes so we can recover on cleanup.
-        original_input_mode = DWORD()
-        if not GetConsoleMode(InHandle, original_input_mode):
-            raise WinError(get_last_error())
-        trace(f'saved original input mode 0x{original_input_mode.value:x}')
-        self.__original_input_mode = original_input_mode.value
+        assert _winapi is not None  # to make mypy happy
+        self.__original_input_mode = _winapi.GetConsoleMode(InHandle)
+        trace(f'saved original input mode 0x{self.__original_input_mode:x}')
 
-        if not SetConsoleMode(
+        _winapi.SetConsoleMode(
             OutHandle,
             ENABLE_WRAP_AT_EOL_OUTPUT
             | ENABLE_PROCESSED_OUTPUT
             | ENABLE_VIRTUAL_TERMINAL_PROCESSING,
-        ):
-            raise WinError(get_last_error())
+        )
 
         self.width = 80
         self.height = 25
@@ -474,8 +471,8 @@ def prepare(self) -> None:
         self.sync_rendered_screen(RenderedScreen.empty(), self.posxy)
 
         if self.__vt_support:
-            if not SetConsoleMode(InHandle, self.__original_input_mode | ENABLE_VIRTUAL_TERMINAL_INPUT):
-                raise WinError(get_last_error())
+            assert _winapi is not None  # to make mypy happy
+            _winapi.SetConsoleMode(InHandle, self.__original_input_mode | ENABLE_VIRTUAL_TERMINAL_INPUT)
             self._enable_bracketed_paste()
 
     def restore(self) -> None:
@@ -484,8 +481,8 @@ def restore(self) -> None:
             self.__write(ANSIColors.RESET)
             # Recover to original mode before running REPL
             self._disable_bracketed_paste()
-            if not SetConsoleMode(InHandle, self.__original_input_mode):
-                raise WinError(get_last_error())
+            assert _winapi is not None  # to make mypy happy
+            _winapi.SetConsoleMode(InHandle, self.__original_input_mode)
 
     def _move_relative(self, x: int, y: int) -> None:
         """Moves relative to the current posxy"""
@@ -661,8 +658,8 @@ def flushoutput(self) -> None:
 
     def forgetinput(self) -> None:
         """Forget all pending, but not yet processed input."""
-        if not FlushConsoleInputBuffer(InHandle):
-            raise WinError(get_last_error())
+        assert _winapi is not None  # to make mypy happy
+        _winapi.FlushConsoleInputBuffer(InHandle)
 
     def getpending(self) -> Event:
         """Return the characters that have been typed but not yet
@@ -815,23 +812,10 @@ class INPUT_RECORD(Structure):
     ]
     ScrollConsoleScreenBuffer.restype = BOOL
 
-    GetConsoleMode = _KERNEL32.GetConsoleMode
-    GetConsoleMode.argtypes = [HANDLE, POINTER(DWORD)]
-    GetConsoleMode.restype = BOOL
-
-    SetConsoleMode = _KERNEL32.SetConsoleMode
-    SetConsoleMode.argtypes = [HANDLE, DWORD]
-    SetConsoleMode.restype = BOOL
-
     ReadConsoleInput = _KERNEL32.ReadConsoleInputW
     ReadConsoleInput.argtypes = [HANDLE, POINTER(INPUT_RECORD), DWORD, POINTER(DWORD)]
     ReadConsoleInput.restype = BOOL
 
-
-    FlushConsoleInputBuffer = _KERNEL32.FlushConsoleInputBuffer
-    FlushConsoleInputBuffer.argtypes = [HANDLE]
-    FlushConsoleInputBuffer.restype = BOOL
-
     OutHandle = _winapi.GetStdHandle(_winapi.STD_OUTPUT_HANDLE)
     InHandle = _winapi.GetStdHandle(_winapi.STD_INPUT_HANDLE)
 else:
@@ -841,9 +825,6 @@ def _win_only(*args, **kwargs):
 
     GetConsoleScreenBufferInfo = _win_only
     ScrollConsoleScreenBuffer = _win_only
-    GetConsoleMode = _win_only
-    SetConsoleMode = _win_only
     ReadConsoleInput = _win_only
-    FlushConsoleInputBuffer = _win_only
     OutHandle = 0
     InHandle = 0

Lib/test/test_pyrepl/test_windows_console.py

@@ -399,7 +399,7 @@ def test_reset_on_restore(self):
         _, con = self.handle_events(events)
         con._WindowsConsole__vt_support = True
         con._WindowsConsole__original_input_mode = 0
-        with patch.object(wc, "SetConsoleMode", return_value=1):
+        with patch.object(wc._winapi, "SetConsoleMode"):
             con.restore()
         con.out.write.assert_any_call(ANSIColors.RESET.encode(con.encoding))
 

Lib/test/test_winapi.py

@@ -166,6 +166,38 @@ def test_namedpipe(self):
             pipe2.flush()
             self.assertEqual((b'testdata', 8), _winapi.PeekNamedPipe(pipe, 8)[:2])
 
+    def test_console_mode_non_console_handle(self):
+        rfd, wfd = _winapi.CreatePipe(None, 0)
+        self.addCleanup(_winapi.CloseHandle, rfd)
+        self.addCleanup(_winapi.CloseHandle, wfd)
+
+        with self.assertRaises(OSError):
+            _winapi.GetConsoleMode(rfd)
+        with self.assertRaises(OSError):
+            _winapi.SetConsoleMode(rfd, 0)
+        with self.assertRaises(OSError):
+            _winapi.FlushConsoleInputBuffer(rfd)
+
+    def test_console_mode(self):
+        try:
+            handle = _winapi.CreateFile(
+                "CONIN$",
+                _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
+                3, # 3=FILE_SHARE_READ|FILE_SHARE_WRITE
+                _winapi.NULL,
+                _winapi.OPEN_EXISTING,
+                0,
+                _winapi.NULL)
+        except OSError:
+            self.skipTest("requires a console")
+        self.addCleanup(_winapi.CloseHandle, handle)
+
+        mode = _winapi.GetConsoleMode(handle)
+        self.assertIsInstance(mode, int)
+        self.assertIsNone(_winapi.SetConsoleMode(handle, mode))
+        self.assertEqual(_winapi.GetConsoleMode(handle), mode)
+        self.assertIsNone(_winapi.FlushConsoleInputBuffer(handle))
+
     def test_event_source_registration(self):
         source_name = "PythonTestEventSource"
 

Modules/_winapi.c

@@ -2824,6 +2824,95 @@ _winapi_GetFileType_impl(PyObject *module, HANDLE handle)
     return result;
 }
 
+/*[clinic input]
+_winapi.GetConsoleMode -> DWORD
+
+    handle: HANDLE
+    /
+
+Return the current mode of a console buffer.
+
+The handle may refer to a console input buffer or a console screen
+buffer.
+[clinic start generated code]*/
+
+static DWORD
+_winapi_GetConsoleMode_impl(PyObject *module, HANDLE handle)
+/*[clinic end generated code: output=4c0155d4357673ac input=4768c2370a960135]*/
+{
+    DWORD mode;
+    BOOL ret;
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = GetConsoleMode(handle, &mode);
+    Py_END_ALLOW_THREADS
+
+    if (!ret) {
+        PyErr_SetFromWindowsErr(0);
+        return PY_DWORD_MAX;
+    }
+
+    return mode;
+}
+
+/*[clinic input]
+_winapi.SetConsoleMode
+
+    handle: HANDLE
+    mode: DWORD(bitwise=True)
+    /
+
+Set the mode of a console buffer.
+
+The handle may refer to a console input buffer or a console screen
+buffer.
+[clinic start generated code]*/
+
+static PyObject *
+_winapi_SetConsoleMode_impl(PyObject *module, HANDLE handle, DWORD mode)
+/*[clinic end generated code: output=3351f12c46a8938f input=41ffc3f65998d7e1]*/
+{
+    BOOL ret;
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = SetConsoleMode(handle, mode);
+    Py_END_ALLOW_THREADS
+
+    if (!ret) {
+        return PyErr_SetFromWindowsErr(0);
+    }
+
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
+_winapi.FlushConsoleInputBuffer
+
+    handle: HANDLE
+    /
+
+Flush the console input buffer.
+
+All input records currently in the input buffer are discarded.
+[clinic start generated code]*/
+
+static PyObject *
+_winapi_FlushConsoleInputBuffer_impl(PyObject *module, HANDLE handle)
+/*[clinic end generated code: output=2edaf920d25233ab input=9684cbfada20c237]*/
+{
+    BOOL ret;
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = FlushConsoleInputBuffer(handle);
+    Py_END_ALLOW_THREADS
+
+    if (!ret) {
+        return PyErr_SetFromWindowsErr(0);
+    }
+
+    Py_RETURN_NONE;
+}
+
 
 /*[clinic input]
 _winapi._mimetypes_read_windows_registry
@@ -3232,6 +3321,9 @@ static PyMethodDef winapi_functions[] = {
     _WINAPI_GETACP_METHODDEF
     _WINAPI_GETOEMCP_METHODDEF
     _WINAPI_GETFILETYPE_METHODDEF
+    _WINAPI_GETCONSOLEMODE_METHODDEF
+    _WINAPI_SETCONSOLEMODE_METHODDEF
+    _WINAPI_FLUSHCONSOLEINPUTBUFFER_METHODDEF
     _WINAPI__MIMETYPES_READ_WINDOWS_REGISTRY_METHODDEF
     _WINAPI_NEEDCURRENTDIRECTORYFOREXEPATH_METHODDEF
     _WINAPI_COPYFILE2_METHODDEF