_winconsole.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370
  1. # -*- coding: utf-8 -*-
  2. # This module is based on the excellent work by Adam Bartoš who
  3. # provided a lot of what went into the implementation here in
  4. # the discussion to issue1602 in the Python bug tracker.
  5. #
  6. # There are some general differences in regards to how this works
  7. # compared to the original patches as we do not need to patch
  8. # the entire interpreter but just work in our little world of
  9. # echo and prmopt.
  10. import ctypes
  11. import io
  12. import os
  13. import sys
  14. import time
  15. import zlib
  16. from ctypes import byref
  17. from ctypes import c_char
  18. from ctypes import c_char_p
  19. from ctypes import c_int
  20. from ctypes import c_ssize_t
  21. from ctypes import c_ulong
  22. from ctypes import c_void_p
  23. from ctypes import POINTER
  24. from ctypes import py_object
  25. from ctypes import windll
  26. from ctypes import WinError
  27. from ctypes import WINFUNCTYPE
  28. from ctypes.wintypes import DWORD
  29. from ctypes.wintypes import HANDLE
  30. from ctypes.wintypes import LPCWSTR
  31. from ctypes.wintypes import LPWSTR
  32. import msvcrt
  33. from ._compat import _NonClosingTextIOWrapper
  34. from ._compat import PY2
  35. from ._compat import text_type
  36. try:
  37. from ctypes import pythonapi
  38. PyObject_GetBuffer = pythonapi.PyObject_GetBuffer
  39. PyBuffer_Release = pythonapi.PyBuffer_Release
  40. except ImportError:
  41. pythonapi = None
  42. c_ssize_p = POINTER(c_ssize_t)
  43. kernel32 = windll.kernel32
  44. GetStdHandle = kernel32.GetStdHandle
  45. ReadConsoleW = kernel32.ReadConsoleW
  46. WriteConsoleW = kernel32.WriteConsoleW
  47. GetConsoleMode = kernel32.GetConsoleMode
  48. GetLastError = kernel32.GetLastError
  49. GetCommandLineW = WINFUNCTYPE(LPWSTR)(("GetCommandLineW", windll.kernel32))
  50. CommandLineToArgvW = WINFUNCTYPE(POINTER(LPWSTR), LPCWSTR, POINTER(c_int))(
  51. ("CommandLineToArgvW", windll.shell32)
  52. )
  53. LocalFree = WINFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p)(
  54. ("LocalFree", windll.kernel32)
  55. )
  56. STDIN_HANDLE = GetStdHandle(-10)
  57. STDOUT_HANDLE = GetStdHandle(-11)
  58. STDERR_HANDLE = GetStdHandle(-12)
  59. PyBUF_SIMPLE = 0
  60. PyBUF_WRITABLE = 1
  61. ERROR_SUCCESS = 0
  62. ERROR_NOT_ENOUGH_MEMORY = 8
  63. ERROR_OPERATION_ABORTED = 995
  64. STDIN_FILENO = 0
  65. STDOUT_FILENO = 1
  66. STDERR_FILENO = 2
  67. EOF = b"\x1a"
  68. MAX_BYTES_WRITTEN = 32767
  69. class Py_buffer(ctypes.Structure):
  70. _fields_ = [
  71. ("buf", c_void_p),
  72. ("obj", py_object),
  73. ("len", c_ssize_t),
  74. ("itemsize", c_ssize_t),
  75. ("readonly", c_int),
  76. ("ndim", c_int),
  77. ("format", c_char_p),
  78. ("shape", c_ssize_p),
  79. ("strides", c_ssize_p),
  80. ("suboffsets", c_ssize_p),
  81. ("internal", c_void_p),
  82. ]
  83. if PY2:
  84. _fields_.insert(-1, ("smalltable", c_ssize_t * 2))
  85. # On PyPy we cannot get buffers so our ability to operate here is
  86. # serverly limited.
  87. if pythonapi is None:
  88. get_buffer = None
  89. else:
  90. def get_buffer(obj, writable=False):
  91. buf = Py_buffer()
  92. flags = PyBUF_WRITABLE if writable else PyBUF_SIMPLE
  93. PyObject_GetBuffer(py_object(obj), byref(buf), flags)
  94. try:
  95. buffer_type = c_char * buf.len
  96. return buffer_type.from_address(buf.buf)
  97. finally:
  98. PyBuffer_Release(byref(buf))
  99. class _WindowsConsoleRawIOBase(io.RawIOBase):
  100. def __init__(self, handle):
  101. self.handle = handle
  102. def isatty(self):
  103. io.RawIOBase.isatty(self)
  104. return True
  105. class _WindowsConsoleReader(_WindowsConsoleRawIOBase):
  106. def readable(self):
  107. return True
  108. def readinto(self, b):
  109. bytes_to_be_read = len(b)
  110. if not bytes_to_be_read:
  111. return 0
  112. elif bytes_to_be_read % 2:
  113. raise ValueError(
  114. "cannot read odd number of bytes from UTF-16-LE encoded console"
  115. )
  116. buffer = get_buffer(b, writable=True)
  117. code_units_to_be_read = bytes_to_be_read // 2
  118. code_units_read = c_ulong()
  119. rv = ReadConsoleW(
  120. HANDLE(self.handle),
  121. buffer,
  122. code_units_to_be_read,
  123. byref(code_units_read),
  124. None,
  125. )
  126. if GetLastError() == ERROR_OPERATION_ABORTED:
  127. # wait for KeyboardInterrupt
  128. time.sleep(0.1)
  129. if not rv:
  130. raise OSError("Windows error: {}".format(GetLastError()))
  131. if buffer[0] == EOF:
  132. return 0
  133. return 2 * code_units_read.value
  134. class _WindowsConsoleWriter(_WindowsConsoleRawIOBase):
  135. def writable(self):
  136. return True
  137. @staticmethod
  138. def _get_error_message(errno):
  139. if errno == ERROR_SUCCESS:
  140. return "ERROR_SUCCESS"
  141. elif errno == ERROR_NOT_ENOUGH_MEMORY:
  142. return "ERROR_NOT_ENOUGH_MEMORY"
  143. return "Windows error {}".format(errno)
  144. def write(self, b):
  145. bytes_to_be_written = len(b)
  146. buf = get_buffer(b)
  147. code_units_to_be_written = min(bytes_to_be_written, MAX_BYTES_WRITTEN) // 2
  148. code_units_written = c_ulong()
  149. WriteConsoleW(
  150. HANDLE(self.handle),
  151. buf,
  152. code_units_to_be_written,
  153. byref(code_units_written),
  154. None,
  155. )
  156. bytes_written = 2 * code_units_written.value
  157. if bytes_written == 0 and bytes_to_be_written > 0:
  158. raise OSError(self._get_error_message(GetLastError()))
  159. return bytes_written
  160. class ConsoleStream(object):
  161. def __init__(self, text_stream, byte_stream):
  162. self._text_stream = text_stream
  163. self.buffer = byte_stream
  164. @property
  165. def name(self):
  166. return self.buffer.name
  167. def write(self, x):
  168. if isinstance(x, text_type):
  169. return self._text_stream.write(x)
  170. try:
  171. self.flush()
  172. except Exception:
  173. pass
  174. return self.buffer.write(x)
  175. def writelines(self, lines):
  176. for line in lines:
  177. self.write(line)
  178. def __getattr__(self, name):
  179. return getattr(self._text_stream, name)
  180. def isatty(self):
  181. return self.buffer.isatty()
  182. def __repr__(self):
  183. return "<ConsoleStream name={!r} encoding={!r}>".format(
  184. self.name, self.encoding
  185. )
  186. class WindowsChunkedWriter(object):
  187. """
  188. Wraps a stream (such as stdout), acting as a transparent proxy for all
  189. attribute access apart from method 'write()' which we wrap to write in
  190. limited chunks due to a Windows limitation on binary console streams.
  191. """
  192. def __init__(self, wrapped):
  193. # double-underscore everything to prevent clashes with names of
  194. # attributes on the wrapped stream object.
  195. self.__wrapped = wrapped
  196. def __getattr__(self, name):
  197. return getattr(self.__wrapped, name)
  198. def write(self, text):
  199. total_to_write = len(text)
  200. written = 0
  201. while written < total_to_write:
  202. to_write = min(total_to_write - written, MAX_BYTES_WRITTEN)
  203. self.__wrapped.write(text[written : written + to_write])
  204. written += to_write
  205. _wrapped_std_streams = set()
  206. def _wrap_std_stream(name):
  207. # Python 2 & Windows 7 and below
  208. if (
  209. PY2
  210. and sys.getwindowsversion()[:2] <= (6, 1)
  211. and name not in _wrapped_std_streams
  212. ):
  213. setattr(sys, name, WindowsChunkedWriter(getattr(sys, name)))
  214. _wrapped_std_streams.add(name)
  215. def _get_text_stdin(buffer_stream):
  216. text_stream = _NonClosingTextIOWrapper(
  217. io.BufferedReader(_WindowsConsoleReader(STDIN_HANDLE)),
  218. "utf-16-le",
  219. "strict",
  220. line_buffering=True,
  221. )
  222. return ConsoleStream(text_stream, buffer_stream)
  223. def _get_text_stdout(buffer_stream):
  224. text_stream = _NonClosingTextIOWrapper(
  225. io.BufferedWriter(_WindowsConsoleWriter(STDOUT_HANDLE)),
  226. "utf-16-le",
  227. "strict",
  228. line_buffering=True,
  229. )
  230. return ConsoleStream(text_stream, buffer_stream)
  231. def _get_text_stderr(buffer_stream):
  232. text_stream = _NonClosingTextIOWrapper(
  233. io.BufferedWriter(_WindowsConsoleWriter(STDERR_HANDLE)),
  234. "utf-16-le",
  235. "strict",
  236. line_buffering=True,
  237. )
  238. return ConsoleStream(text_stream, buffer_stream)
  239. if PY2:
  240. def _hash_py_argv():
  241. return zlib.crc32("\x00".join(sys.argv[1:]))
  242. _initial_argv_hash = _hash_py_argv()
  243. def _get_windows_argv():
  244. argc = c_int(0)
  245. argv_unicode = CommandLineToArgvW(GetCommandLineW(), byref(argc))
  246. if not argv_unicode:
  247. raise WinError()
  248. try:
  249. argv = [argv_unicode[i] for i in range(0, argc.value)]
  250. finally:
  251. LocalFree(argv_unicode)
  252. del argv_unicode
  253. if not hasattr(sys, "frozen"):
  254. argv = argv[1:]
  255. while len(argv) > 0:
  256. arg = argv[0]
  257. if not arg.startswith("-") or arg == "-":
  258. break
  259. argv = argv[1:]
  260. if arg.startswith(("-c", "-m")):
  261. break
  262. return argv[1:]
  263. _stream_factories = {
  264. 0: _get_text_stdin,
  265. 1: _get_text_stdout,
  266. 2: _get_text_stderr,
  267. }
  268. def _is_console(f):
  269. if not hasattr(f, "fileno"):
  270. return False
  271. try:
  272. fileno = f.fileno()
  273. except OSError:
  274. return False
  275. handle = msvcrt.get_osfhandle(fileno)
  276. return bool(GetConsoleMode(handle, byref(DWORD())))
  277. def _get_windows_console_stream(f, encoding, errors):
  278. if (
  279. get_buffer is not None
  280. and encoding in ("utf-16-le", None)
  281. and errors in ("strict", None)
  282. and _is_console(f)
  283. ):
  284. func = _stream_factories.get(f.fileno())
  285. if func is not None:
  286. if not PY2:
  287. f = getattr(f, "buffer", None)
  288. if f is None:
  289. return None
  290. else:
  291. # If we are on Python 2 we need to set the stream that we
  292. # deal with to binary mode as otherwise the exercise if a
  293. # bit moot. The same problems apply as for
  294. # get_binary_stdin and friends from _compat.
  295. msvcrt.setmode(f.fileno(), os.O_BINARY)
  296. return func(f)