_winconsole.py 8.8 KB

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