fix_encoding.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324
  1. # Copyright (c) 2011 The Chromium Authors. All rights reserved.
  2. # Use of this source code is governed by a BSD-style license that can be
  3. # found in the LICENSE file.
  4. """Collection of functions and classes to fix various encoding problems on
  5. multiple platforms with python.
  6. """
  7. from __future__ import print_function
  8. import codecs
  9. import locale
  10. import os
  11. import sys
  12. def complain(message):
  13. """If any exception occurs in this file, we'll probably try to print it
  14. on stderr, which makes for frustrating debugging if stderr is directed
  15. to our wrapper. So be paranoid about catching errors and reporting them
  16. to sys.__stderr__, so that the user has a higher chance to see them.
  17. """
  18. print(isinstance(message, str) and message or repr(message),
  19. file=sys.__stderr__)
  20. def fix_default_encoding():
  21. """Forces utf8 solidly on all platforms.
  22. By default python execution environment is lazy and defaults to ascii
  23. encoding.
  24. http://uucode.com/blog/2007/03/23/shut-up-you-dummy-7-bit-python/
  25. """
  26. if sys.getdefaultencoding() == 'utf-8':
  27. return False
  28. # Regenerate setdefaultencoding.
  29. reload(sys)
  30. # Module 'sys' has no 'setdefaultencoding' member
  31. # pylint: disable=no-member
  32. sys.setdefaultencoding('utf-8')
  33. for attr in dir(locale):
  34. if attr[0:3] != 'LC_':
  35. continue
  36. aref = getattr(locale, attr)
  37. try:
  38. locale.setlocale(aref, '')
  39. except locale.Error:
  40. continue
  41. try:
  42. lang, _ = locale.getdefaultlocale()
  43. except (TypeError, ValueError):
  44. continue
  45. if lang:
  46. try:
  47. locale.setlocale(aref, (lang, 'UTF-8'))
  48. except locale.Error:
  49. os.environ[attr] = lang + '.UTF-8'
  50. try:
  51. locale.setlocale(locale.LC_ALL, '')
  52. except locale.Error:
  53. pass
  54. return True
  55. ###############################
  56. # Windows specific
  57. def fix_win_codec():
  58. """Works around <http://bugs.python.org/issue6058>."""
  59. # <http://msdn.microsoft.com/en-us/library/dd317756.aspx>
  60. try:
  61. codecs.lookup('cp65001')
  62. return False
  63. except LookupError:
  64. codecs.register(
  65. lambda name: name == 'cp65001' and codecs.lookup('utf-8') or None)
  66. return True
  67. class WinUnicodeOutputBase(object):
  68. """Base class to adapt sys.stdout or sys.stderr to behave correctly on
  69. Windows.
  70. Setting encoding to utf-8 is recommended.
  71. """
  72. def __init__(self, fileno, name, encoding):
  73. # Corresponding file handle.
  74. self._fileno = fileno
  75. self.encoding = encoding
  76. self.name = name
  77. self.closed = False
  78. self.softspace = False
  79. self.mode = 'w'
  80. @staticmethod
  81. def isatty():
  82. return False
  83. def close(self):
  84. # Don't really close the handle, that would only cause problems.
  85. self.closed = True
  86. def fileno(self):
  87. return self._fileno
  88. def flush(self):
  89. raise NotImplementedError()
  90. def write(self, text):
  91. raise NotImplementedError()
  92. def writelines(self, lines):
  93. try:
  94. for line in lines:
  95. self.write(line)
  96. except Exception as e:
  97. complain('%s.writelines: %r' % (self.name, e))
  98. raise
  99. class WinUnicodeConsoleOutput(WinUnicodeOutputBase):
  100. """Output adapter to a Windows Console.
  101. Understands how to use the win32 console API.
  102. """
  103. def __init__(self, console_handle, fileno, stream_name, encoding):
  104. super(WinUnicodeConsoleOutput,
  105. self).__init__(fileno, '<Unicode console %s>' % stream_name,
  106. encoding)
  107. # Handle to use for WriteConsoleW
  108. self._console_handle = console_handle
  109. # Loads the necessary function.
  110. # These types are available on linux but not Mac.
  111. # pylint: disable=no-name-in-module,F0401
  112. from ctypes import byref, GetLastError, POINTER, windll, WINFUNCTYPE
  113. from ctypes.wintypes import BOOL, DWORD, HANDLE, LPWSTR
  114. from ctypes.wintypes import LPVOID # pylint: disable=no-name-in-module
  115. self._DWORD = DWORD
  116. self._byref = byref
  117. # <http://msdn.microsoft.com/en-us/library/ms687401.aspx>
  118. self._WriteConsoleW = WINFUNCTYPE(BOOL, HANDLE, LPWSTR, DWORD,
  119. POINTER(DWORD),
  120. LPVOID)(('WriteConsoleW',
  121. windll.kernel32))
  122. self._GetLastError = GetLastError
  123. def flush(self):
  124. # No need to flush the console since it's immediate.
  125. pass
  126. def write(self, text):
  127. try:
  128. if isinstance(text, bytes):
  129. # Bytestrings need to be decoded to a string before being passed
  130. # to Windows.
  131. text = text.decode(self.encoding, 'replace')
  132. remaining = len(text)
  133. while remaining > 0:
  134. n = self._DWORD(0)
  135. # There is a shorter-than-documented limitation on the length of
  136. # the string passed to WriteConsoleW. See
  137. # <http://tahoe-lafs.org/trac/tahoe-lafs/ticket/1232>.
  138. retval = self._WriteConsoleW(self._console_handle, text,
  139. min(remaining, 10000),
  140. self._byref(n), None)
  141. if retval == 0 or n.value == 0:
  142. raise IOError('WriteConsoleW returned %r, n.value = %r, '
  143. 'last error = %r' %
  144. (retval, n.value, self._GetLastError()))
  145. remaining -= n.value
  146. if not remaining:
  147. break
  148. text = text[int(n.value):]
  149. except Exception as e:
  150. complain('%s.write: %r' % (self.name, e))
  151. raise
  152. class WinUnicodeOutput(WinUnicodeOutputBase):
  153. """Output adaptor to a file output on Windows.
  154. If the standard FileWrite function is used, it will be encoded in the current
  155. code page. WriteConsoleW() permits writing any character.
  156. """
  157. def __init__(self, stream, fileno, encoding):
  158. super(WinUnicodeOutput,
  159. self).__init__(fileno, '<Unicode redirected %s>' % stream.name,
  160. encoding)
  161. # Output stream
  162. self._stream = stream
  163. # Flush right now.
  164. self.flush()
  165. def flush(self):
  166. try:
  167. self._stream.flush()
  168. except Exception as e:
  169. complain('%s.flush: %r from %r' % (self.name, e, self._stream))
  170. raise
  171. def write(self, text):
  172. try:
  173. if isinstance(text, bytes):
  174. # Replace characters that cannot be printed instead of failing.
  175. text = text.decode(self.encoding, 'replace')
  176. # When redirecting to a file or process any \n characters will be
  177. # replaced with \r\n. If the text to be printed already has \r\n
  178. # line endings then \r\r\n line endings will be generated, leading
  179. # to double-spacing of some output. Normalizing line endings to \n
  180. # avoids this problem.
  181. text = text.replace('\r\n', '\n')
  182. self._stream.write(text)
  183. except Exception as e:
  184. complain('%s.write: %r' % (self.name, e))
  185. raise
  186. def win_handle_is_a_console(handle):
  187. """Returns True if a Windows file handle is a handle to a console."""
  188. # These types are available on linux but not Mac.
  189. # pylint: disable=no-name-in-module,F0401
  190. from ctypes import byref, POINTER, windll, WINFUNCTYPE
  191. from ctypes.wintypes import BOOL, DWORD, HANDLE
  192. FILE_TYPE_CHAR = 0x0002
  193. FILE_TYPE_REMOTE = 0x8000
  194. INVALID_HANDLE_VALUE = DWORD(-1).value
  195. # <http://msdn.microsoft.com/en-us/library/ms683167.aspx>
  196. GetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, POINTER(DWORD))(
  197. ('GetConsoleMode', windll.kernel32))
  198. # <http://msdn.microsoft.com/en-us/library/aa364960.aspx>
  199. GetFileType = WINFUNCTYPE(DWORD, DWORD)(('GetFileType', windll.kernel32))
  200. # GetStdHandle returns INVALID_HANDLE_VALUE, NULL, or a valid handle.
  201. if handle == INVALID_HANDLE_VALUE or handle is None:
  202. return False
  203. return ((GetFileType(handle) & ~FILE_TYPE_REMOTE) == FILE_TYPE_CHAR
  204. and GetConsoleMode(handle, byref(DWORD())))
  205. def win_get_unicode_stream(stream, excepted_fileno, output_handle, encoding):
  206. """Returns a unicode-compatible stream.
  207. This function will return a direct-Console writing object only if:
  208. - the file number is the expected console file number
  209. - the handle the expected file handle
  210. - the 'real' handle is in fact a handle to a console.
  211. """
  212. old_fileno = getattr(stream, 'fileno', lambda: None)()
  213. if old_fileno == excepted_fileno:
  214. # These types are available on linux but not Mac.
  215. # pylint: disable=no-name-in-module,F0401
  216. from ctypes import windll, WINFUNCTYPE
  217. from ctypes.wintypes import DWORD, HANDLE
  218. # <http://msdn.microsoft.com/en-us/library/ms683231.aspx>
  219. GetStdHandle = WINFUNCTYPE(HANDLE,
  220. DWORD)(('GetStdHandle', windll.kernel32))
  221. real_output_handle = GetStdHandle(DWORD(output_handle))
  222. if win_handle_is_a_console(real_output_handle):
  223. # It's a console.
  224. return WinUnicodeConsoleOutput(real_output_handle, old_fileno,
  225. stream.name, encoding)
  226. # It's something else. Create an auto-encoding stream.
  227. return WinUnicodeOutput(stream, old_fileno, encoding)
  228. def fix_win_console(encoding):
  229. """Makes Unicode console output work independently of the current code page.
  230. This also fixes <http://bugs.python.org/issue1602>.
  231. Credit to Michael Kaplan
  232. <http://blogs.msdn.com/b/michkap/archive/2010/04/07/9989346.aspx> and
  233. TZOmegaTZIOY
  234. <http://stackoverflow.com/questions/878972/windows-cmd-encoding-change-causes-python-crash/1432462#1432462>.
  235. """
  236. if (isinstance(sys.stdout, WinUnicodeOutputBase)
  237. or isinstance(sys.stderr, WinUnicodeOutputBase)):
  238. return False
  239. try:
  240. # SetConsoleCP and SetConsoleOutputCP could be used to change the code
  241. # page but it's not really useful since the code here is using
  242. # WriteConsoleW(). Also, changing the code page is 'permanent' to the
  243. # console and needs to be reverted manually. In practice one needs to
  244. # set the console font to a TTF font to be able to see all the
  245. # characters but it failed for me in practice. In any case, it won't
  246. # throw any exception when printing, which is the important part. -11
  247. # and -12 are defined in stdio.h
  248. sys.stdout = win_get_unicode_stream(sys.stdout, 1, -11, encoding)
  249. sys.stderr = win_get_unicode_stream(sys.stderr, 2, -12, encoding)
  250. # TODO(maruel): Do sys.stdin with ReadConsoleW(). Albeit the limitation
  251. # is "It doesn't appear to be possible to read Unicode characters in
  252. # UTF-8 mode" and this appears to be a limitation of cmd.exe.
  253. except Exception as e:
  254. complain('exception %r while fixing up sys.stdout and sys.stderr' % e)
  255. return True
  256. def fix_encoding():
  257. """Fixes various encoding problems on all platforms.
  258. Should be called at the very beginning of the process.
  259. """
  260. ret = True
  261. if sys.platform == 'win32':
  262. ret &= fix_win_codec()
  263. ret &= fix_default_encoding()
  264. if sys.platform == 'win32':
  265. encoding = sys.getdefaultencoding()
  266. ret &= fix_win_console(encoding)
  267. return ret