Coverage for typer / _click / _winconsole.py: 100%

1# This module is based on the excellent work by Adam Bartoš who 

2# provided a lot of what went into the implementation here in 

3# the discussion to issue1602 in the Python bug tracker. 

4# 

5# There are some general differences in regards to how this works 

6# compared to the original patches as we do not need to patch 

7# the entire interpreter but just work in our little world of 

8# echo and prompt. 

9import io 2 ctx1ab

10import sys 2 ctx1ab

11import time 2 ctx1ab

12from collections.abc import Callable, Iterable, Mapping 2 ctx1ab

13from ctypes import ( 2 ctx1ab

14 POINTER, 

15 Array, 

16 Structure, 

17 byref, 

18 c_char, 

19 c_char_p, 

20 c_int, 

21 c_ssize_t, 

22 c_ulong, 

23 c_void_p, 

24 py_object, 

25) 

26from ctypes.wintypes import DWORD, HANDLE, LPCWSTR, LPWSTR 2 ctx1ab

27from typing import ( 2 ctx1ab

28 TYPE_CHECKING, 

29 Any, 

30 AnyStr, 

31 BinaryIO, 

32 Literal, 

33 TextIO, 

34 cast, 

35) 

36 

37from ._compat import _NonClosingTextIOWrapper 2 ctx1ab

38 

39assert sys.platform == "win32" 2 ctx1ab

40import msvcrt # noqa: E402 2 ctx1ab

41from ctypes import WINFUNCTYPE, windll # noqa: E402 2 ctx1ab

42 

43c_ssize_p = POINTER(c_ssize_t) 2 ctx1ab

44 

45kernel32 = windll.kernel32 2 ctx1ab

46GetStdHandle = kernel32.GetStdHandle 2 ctx1ab

47ReadConsoleW = kernel32.ReadConsoleW 2 ctx1ab

48WriteConsoleW = kernel32.WriteConsoleW 2 ctx1ab

49GetConsoleMode = kernel32.GetConsoleMode 2 ctx1ab

50GetLastError = kernel32.GetLastError 2 ctx1ab

51GetCommandLineW = WINFUNCTYPE(LPWSTR)(("GetCommandLineW", windll.kernel32)) 2 ctx1ab

52CommandLineToArgvW = WINFUNCTYPE(POINTER(LPWSTR), LPCWSTR, POINTER(c_int))( 2 ctx1ab

53 ("CommandLineToArgvW", windll.shell32) 

54) 

55LocalFree = WINFUNCTYPE(c_void_p, c_void_p)(("LocalFree", windll.kernel32)) 2 ctx1ab

56 

57STDIN_HANDLE = GetStdHandle(-10) 2 ctx1ab

58STDOUT_HANDLE = GetStdHandle(-11) 2 ctx1ab

59STDERR_HANDLE = GetStdHandle(-12) 2 ctx1ab

60 

61PyBUF_SIMPLE = 0 2 ctx1ab

62PyBUF_WRITABLE = 1 2 ctx1ab

63 

64ERROR_SUCCESS = 0 2 ctx1ab

65ERROR_NOT_ENOUGH_MEMORY = 8 2 ctx1ab

66ERROR_OPERATION_ABORTED = 995 2 ctx1ab

67 

68STDIN_FILENO = 0 2 ctx1ab

69STDOUT_FILENO = 1 2 ctx1ab

70STDERR_FILENO = 2 2 ctx1ab

71 

72EOF = b"\x1a" 2 ctx1ab

73MAX_BYTES_WRITTEN = 32767 2 ctx1ab

74 

75if TYPE_CHECKING: 2 ctx1ab

76 try: 

77 # Using `typing_extensions.Buffer` instead of `collections.abc` 

78 # on Windows for some reason does not have `Sized` implemented. 

79 from collections.abc import Buffer # type: ignore 

80 except ImportError: 

81 from typing_extensions import Buffer 

82 

83try: 2 ctx1ab

84 from ctypes import pythonapi 2 ctx1ab

85except ImportError: # pragma: no cover 

86 # On PyPy we cannot get buffers so our ability to operate here is 

87 # severely limited. 

88 get_buffer = None 

89else: 

90 

91 class Py_buffer(Structure): 2 ctx1ab

92 _fields_ = [ # noqa: RUF012 2 ctx1ab

93 ("buf", c_void_p), 

94 ("obj", py_object), 

95 ("len", c_ssize_t), 

96 ("itemsize", c_ssize_t), 

97 ("readonly", c_int), 

98 ("ndim", c_int), 

99 ("format", c_char_p), 

100 ("shape", c_ssize_p), 

101 ("strides", c_ssize_p), 

102 ("suboffsets", c_ssize_p), 

103 ("internal", c_void_p), 

104 ] 

105 

106 PyObject_GetBuffer = pythonapi.PyObject_GetBuffer 2 ctx1ab

107 PyBuffer_Release = pythonapi.PyBuffer_Release 2 ctx1ab

108 

109 def get_buffer(obj: "Buffer", writable: bool = False) -> Array[c_char]: 2 ctx1ab

110 buf = Py_buffer() 2 ctx1ab

111 flags: int = PyBUF_WRITABLE if writable else PyBUF_SIMPLE 2 ctx1ab

112 PyObject_GetBuffer(py_object(obj), byref(buf), flags) 2 ctx1ab

113 

114 try: 2 ctx1ab

115 buffer_type = c_char * buf.len 2 ctx1ab

116 out: Array[c_char] = buffer_type.from_address(buf.buf) 2 ctx1ab

117 return out 2 ctx1ab

118 finally: 

119 PyBuffer_Release(byref(buf)) 2 ctx1ab

120 

121 

122class _WindowsConsoleRawIOBase(io.RawIOBase): 2 ctx1ab

123 def __init__(self, handle: int | None) -> None: 2 ctx1ab

124 self.handle = handle 2 ctx1ab

125 

126 def isatty(self) -> Literal[True]: 2 ctx1ab

127 super().isatty() 2 ctx1ab

128 return True 2 ctx1ab

129 

130 

131class _WindowsConsoleReader(_WindowsConsoleRawIOBase): 2 ctx1ab

132 def readable(self) -> Literal[True]: 2 ctx1ab

133 return True 2 ctx1ab

134 

135 def readinto(self, b: "Buffer") -> int: 2 ctx1ab

136 bytes_to_be_read = len(b) 2 ctx1ab

137 if not bytes_to_be_read: 2 ctx1ab

138 return 0 2 ctx1ab

139 elif bytes_to_be_read % 2: 2 ctx1ab

140 raise ValueError( 2 ctx1ab

141 "cannot read odd number of bytes from UTF-16-LE encoded console" 

142 ) 

143 

144 buffer = get_buffer(b, writable=True) 2 ctx1ab

145 code_units_to_be_read = bytes_to_be_read // 2 2 ctx1ab

146 code_units_read = c_ulong() 2 ctx1ab

147 

148 rv = ReadConsoleW( 2 ctx1ab

149 HANDLE(self.handle), 

150 buffer, 

151 code_units_to_be_read, 

152 byref(code_units_read), 

153 None, 

154 ) 

155 if GetLastError() == ERROR_OPERATION_ABORTED: 2 ctx1ab

156 # wait for KeyboardInterrupt 

157 time.sleep(0.1) 2 ctx1ab

158 if not rv: 2 ctx1ab

159 raise OSError(f"Windows error: {GetLastError()}") 2 ctx1ab

160 

161 if buffer[0] == EOF: 2 ctx1ab

162 return 0 2 ctx1ab

163 return 2 * code_units_read.value 2 ctx1ab

164 

165 

166class _WindowsConsoleWriter(_WindowsConsoleRawIOBase): 2 ctx1ab

167 def writable(self) -> Literal[True]: 2 ctx1ab

168 return True 2 ctx1ab

169 

170 @staticmethod 2 ctx1ab

171 def _get_error_message(errno: int) -> str: 2 ctx1ab

172 if errno == ERROR_SUCCESS: 2 ctx1ab

173 return "ERROR_SUCCESS" 2 ctx1ab

174 elif errno == ERROR_NOT_ENOUGH_MEMORY: 2 ctx1ab

175 return "ERROR_NOT_ENOUGH_MEMORY" 2 ctx1ab

176 return f"Windows error {errno}" 2 ctx1ab

177 

178 def write(self, b: "Buffer") -> int: 2 ctx1ab

179 bytes_to_be_written = len(b) 2 ctx1ab

180 buf = get_buffer(b) 2 ctx1ab

181 code_units_to_be_written = min(bytes_to_be_written, MAX_BYTES_WRITTEN) // 2 2 ctx1ab

182 code_units_written = c_ulong() 2 ctx1ab

183 

184 WriteConsoleW( 2 ctx1ab

185 HANDLE(self.handle), 

186 buf, 

187 code_units_to_be_written, 

188 byref(code_units_written), 

189 None, 

190 ) 

191 bytes_written = 2 * code_units_written.value 2 ctx1ab

192 

193 if bytes_written == 0 and bytes_to_be_written > 0: 2 ctx1ab

194 raise OSError(self._get_error_message(GetLastError())) # pragma: no cover 

195 return bytes_written 2 ctx1ab

196 

197 

198class ConsoleStream: 2 ctx1ab

199 def __init__(self, text_stream: TextIO, byte_stream: BinaryIO) -> None: 2 ctx1ab

200 self._text_stream = text_stream 2 ctx1ab

201 self.buffer = byte_stream 2 ctx1ab

202 

203 @property 2 ctx1ab

204 def name(self) -> str: 2 ctx1ab

205 return self.buffer.name 2 ctx1ab

206 

207 def write(self, x: AnyStr) -> int: 2 ctx1ab

208 if isinstance(x, str): 2 ctx1ab

209 return self._text_stream.write(x) 2 ctx1ab

210 try: 2 ctx1ab

211 self.flush() 2 ctx1ab

212 except Exception: # pragma: no cover 

213 pass 

214 return self.buffer.write(x) 2 ctx1ab

215 

216 def writelines(self, lines: Iterable[AnyStr]) -> None: 2 ctx1ab

217 for line in lines: 2 ctx1ab

218 self.write(line) 2 ctx1ab

219 

220 def __getattr__(self, name: str) -> Any: 2 ctx1ab

221 return getattr(self._text_stream, name) 2 ctx1ab

222 

223 def isatty(self) -> bool: 2 ctx1ab

224 return self.buffer.isatty() 2 ctx1ab

225 

226 def __repr__(self) -> str: 2 ctx1ab

227 return f"<ConsoleStream name={self.name!r} encoding={self.encoding!r}>" 2 ctx1ab

228 

229 

230def _get_text_stdin(buffer_stream: BinaryIO) -> TextIO: 2 ctx1ab

231 text_stream = _NonClosingTextIOWrapper( 2 ctx1ab

232 io.BufferedReader(_WindowsConsoleReader(STDIN_HANDLE)), 

233 "utf-16-le", 

234 "strict", 

235 line_buffering=True, 

236 ) 

237 return cast(TextIO, ConsoleStream(text_stream, buffer_stream)) 2 ctx1ab

238 

239 

240def _get_text_stdout(buffer_stream: BinaryIO) -> TextIO: 2 ctx1ab

241 text_stream = _NonClosingTextIOWrapper( 2 ctx1ab

242 io.BufferedWriter(_WindowsConsoleWriter(STDOUT_HANDLE)), 

243 "utf-16-le", 

244 "strict", 

245 line_buffering=True, 

246 ) 

247 return cast(TextIO, ConsoleStream(text_stream, buffer_stream)) 2 ctx1ab

248 

249 

250def _get_text_stderr(buffer_stream: BinaryIO) -> TextIO: 2 ctx1ab

251 text_stream = _NonClosingTextIOWrapper( 2 ctx1ab

252 io.BufferedWriter(_WindowsConsoleWriter(STDERR_HANDLE)), 

253 "utf-16-le", 

254 "strict", 

255 line_buffering=True, 

256 ) 

257 return cast(TextIO, ConsoleStream(text_stream, buffer_stream)) 2 ctx1ab

258 

259 

260_stream_factories: Mapping[int, Callable[[BinaryIO], TextIO]] = { 2 ctx1ab

261 0: _get_text_stdin, 

262 1: _get_text_stdout, 

263 2: _get_text_stderr, 

264} 

265 

266 

267def _is_console(f: TextIO) -> bool: 2 ctx1ab

268 if not hasattr(f, "fileno"): 2 ctx1ab

269 return False 2 ctx1ab

270 

271 try: 2 ctx1ab

272 fileno = f.fileno() 2 ctx1ab

273 except (OSError, io.UnsupportedOperation): 2 ctx1ab

274 return False 2 ctx1ab

275 

276 handle = msvcrt.get_osfhandle(fileno) 2 ctx1ab

277 return bool(GetConsoleMode(handle, byref(DWORD()))) 2 ctx1ab

278 

279 

280def _get_windows_console_stream( 2 ctx1ab

281 f: TextIO, encoding: str | None, errors: str | None 

282) -> TextIO | None: 

283 if ( 

284 get_buffer is None 

285 or encoding not in {"utf-16-le", None} 

286 or errors not in {"strict", None} 

287 or not _is_console(f) 

288 ): 

289 return None 2 ctx1ab

290 

291 func = _stream_factories.get(f.fileno()) 2 ctx1ab

292 if func is None: 2 ctx1ab

293 return None 2 ctx1ab

294 

295 b = getattr(f, "buffer", None) 2 ctx1ab

296 

297 if b is None: 2 ctx1ab

298 return None 2 ctx1ab

299 

300 return func(b) 2 ctx1ab