fix_encoding.py 11.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42
# Copyright (c) 2011 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Collection of functions and classes to fix various encoding problems on
multiple platforms with python.
"""

import codecs
import locale
import os
import sys


# Prevents initializing multiple times.
_SYS_ARGV_PROCESSED = False


def complain(message):
  """If any exception occurs in this file, we'll probably try to print it
  on stderr, which makes for frustrating debugging if stderr is directed
  to our wrapper. So be paranoid about catching errors and reporting them
  to sys.__stderr__, so that the user has a higher chance to see them.
  """
  print >> sys.__stderr__, (
      isinstance(message, str) and message or repr(message))


def fix_default_encoding():
  """Forces utf8 solidly on all platforms.

  By default python execution environment is lazy and defaults to ascii
  encoding.

  http://uucode.com/blog/2007/03/23/shut-up-you-dummy-7-bit-python/
  """
  if sys.getdefaultencoding() == 'utf-8':
    return False

  # Regenerate setdefaultencoding.
  reload(sys)
  # Module 'sys' has no 'setdefaultencoding' member
43
  # pylint: disable=no-member
44 45 46 47 48
  sys.setdefaultencoding('utf-8')
  for attr in dir(locale):
    if attr[0:3] != 'LC_':
      continue
    aref = getattr(locale, attr)
49 50 51 52
    try:
      locale.setlocale(aref, '')
    except locale.Error:
      continue
53
    try:
54
      lang, _ = locale.getdefaultlocale()
55
    except (TypeError, ValueError):
56
      continue
57 58 59 60 61
    if lang:
      try:
        locale.setlocale(aref, (lang, 'UTF-8'))
      except locale.Error:
        os.environ[attr] = lang + '.UTF-8'
62 63 64 65
  try:
    locale.setlocale(locale.LC_ALL, '')
  except locale.Error:
    pass
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
  return True


###############################
# Windows specific


def fix_win_sys_argv(encoding):
  """Converts sys.argv to 'encoding' encoded string.

  utf-8 is recommended.

  Works around <http://bugs.python.org/issue2128>.
  """
  global _SYS_ARGV_PROCESSED
  if _SYS_ARGV_PROCESSED:
    return False

84
  # These types are available on linux but not Mac.
85
  # pylint: disable=no-name-in-module,F0401
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173
  from ctypes import byref, c_int, POINTER, windll, WINFUNCTYPE
  from ctypes.wintypes import LPCWSTR, LPWSTR

  # <http://msdn.microsoft.com/en-us/library/ms683156.aspx>
  GetCommandLineW = WINFUNCTYPE(LPWSTR)(('GetCommandLineW', windll.kernel32))
  # <http://msdn.microsoft.com/en-us/library/bb776391.aspx>
  CommandLineToArgvW = WINFUNCTYPE(POINTER(LPWSTR), LPCWSTR, POINTER(c_int))(
      ('CommandLineToArgvW', windll.shell32))

  argc = c_int(0)
  argv_unicode = CommandLineToArgvW(GetCommandLineW(), byref(argc))
  argv = [
      argv_unicode[i].encode(encoding, 'replace')
      for i in xrange(0, argc.value)]

  if not hasattr(sys, 'frozen'):
    # If this is an executable produced by py2exe or bbfreeze, then it
    # will have been invoked directly. Otherwise, unicode_argv[0] is the
    # Python interpreter, so skip that.
    argv = argv[1:]

    # Also skip option arguments to the Python interpreter.
    while len(argv) > 0:
      arg = argv[0]
      if not arg.startswith(u'-') or arg == u'-':
        break
      argv = argv[1:]
      if arg == u'-m':
        # sys.argv[0] should really be the absolute path of the
        # module source, but never mind.
        break
      if arg == u'-c':
        argv[0] = u'-c'
        break
  sys.argv = argv
  _SYS_ARGV_PROCESSED = True
  return True


def fix_win_codec():
  """Works around <http://bugs.python.org/issue6058>."""
  # <http://msdn.microsoft.com/en-us/library/dd317756.aspx>
  try:
    codecs.lookup('cp65001')
    return False
  except LookupError:
    codecs.register(
        lambda name: name == 'cp65001' and codecs.lookup('utf-8') or None)
    return True


class WinUnicodeOutputBase(object):
  """Base class to adapt sys.stdout or sys.stderr to behave correctly on
  Windows.

  Setting encoding to utf-8 is recommended.
  """
  def __init__(self, fileno, name, encoding):
    # Corresponding file handle.
    self._fileno = fileno
    self.encoding = encoding
    self.name = name

    self.closed = False
    self.softspace = False
    self.mode = 'w'

  @staticmethod
  def isatty():
    return False

  def close(self):
    # Don't really close the handle, that would only cause problems.
    self.closed = True

  def fileno(self):
    return self._fileno

  def flush(self):
    raise NotImplementedError()

  def write(self, text):
    raise NotImplementedError()

  def writelines(self, lines):
    try:
      for line in lines:
        self.write(line)
174
    except Exception as e:
175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190
      complain('%s.writelines: %r' % (self.name, e))
      raise


class WinUnicodeConsoleOutput(WinUnicodeOutputBase):
  """Output adapter to a Windows Console.

  Understands how to use the win32 console API.
  """
  def __init__(self, console_handle, fileno, stream_name, encoding):
    super(WinUnicodeConsoleOutput, self).__init__(
        fileno, '<Unicode console %s>' % stream_name, encoding)
    # Handle to use for WriteConsoleW
    self._console_handle = console_handle

    # Loads the necessary function.
191
    # These types are available on linux but not Mac.
192
    # pylint: disable=no-name-in-module,F0401
193
    from ctypes import byref, GetLastError, POINTER, windll, WINFUNCTYPE
194
    from ctypes.wintypes import BOOL, DWORD, HANDLE, LPWSTR
195
    from ctypes.wintypes import LPVOID  # pylint: disable=no-name-in-module
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231

    self._DWORD = DWORD
    self._byref = byref

    # <http://msdn.microsoft.com/en-us/library/ms687401.aspx>
    self._WriteConsoleW = WINFUNCTYPE(
        BOOL, HANDLE, LPWSTR, DWORD, POINTER(DWORD), LPVOID)(
            ('WriteConsoleW', windll.kernel32))
    self._GetLastError = GetLastError

  def flush(self):
    # No need to flush the console since it's immediate.
    pass

  def write(self, text):
    try:
      if not isinstance(text, unicode):
        # Convert to unicode.
        text = str(text).decode(self.encoding, 'replace')
      remaining = len(text)
      while remaining > 0:
        n = self._DWORD(0)
        # There is a shorter-than-documented limitation on the length of the
        # string passed to WriteConsoleW. See
        # <http://tahoe-lafs.org/trac/tahoe-lafs/ticket/1232>.
        retval = self._WriteConsoleW(
            self._console_handle, text,
            min(remaining, 10000),
            self._byref(n), None)
        if retval == 0 or n.value == 0:
          raise IOError(
              'WriteConsoleW returned %r, n.value = %r, last error = %r' % (
                retval, n.value, self._GetLastError()))
        remaining -= n.value
        if not remaining:
          break
232
        text = text[int(n.value):]
233
    except Exception as e:
234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255
      complain('%s.write: %r' % (self.name, e))
      raise


class WinUnicodeOutput(WinUnicodeOutputBase):
  """Output adaptor to a file output on Windows.

  If the standard FileWrite function is used, it will be encoded in the current
  code page. WriteConsoleW() permits writting any character.
  """
  def __init__(self, stream, fileno, encoding):
    super(WinUnicodeOutput, self).__init__(
        fileno, '<Unicode redirected %s>' % stream.name, encoding)
    # Output stream
    self._stream = stream

    # Flush right now.
    self.flush()

  def flush(self):
    try:
      self._stream.flush()
256
    except Exception as e:
257 258 259 260 261 262 263 264 265
      complain('%s.flush: %r from %r' % (self.name, e, self._stream))
      raise

  def write(self, text):
    try:
      if isinstance(text, unicode):
        # Replace characters that cannot be printed instead of failing.
        text = text.encode(self.encoding, 'replace')
      self._stream.write(text)
266
    except Exception as e:
267 268 269 270 271 272
      complain('%s.write: %r' % (self.name, e))
      raise


def win_handle_is_a_console(handle):
  """Returns True if a Windows file handle is a handle to a console."""
273
  # These types are available on linux but not Mac.
274
  # pylint: disable=no-name-in-module,F0401
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305
  from ctypes import byref, POINTER, windll, WINFUNCTYPE
  from ctypes.wintypes import BOOL, DWORD, HANDLE

  FILE_TYPE_CHAR   = 0x0002
  FILE_TYPE_REMOTE = 0x8000
  INVALID_HANDLE_VALUE = DWORD(-1).value

  # <http://msdn.microsoft.com/en-us/library/ms683167.aspx>
  GetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, POINTER(DWORD))(
      ('GetConsoleMode', windll.kernel32))
  # <http://msdn.microsoft.com/en-us/library/aa364960.aspx>
  GetFileType = WINFUNCTYPE(DWORD, DWORD)(('GetFileType', windll.kernel32))

  # GetStdHandle returns INVALID_HANDLE_VALUE, NULL, or a valid handle.
  if handle == INVALID_HANDLE_VALUE or handle is None:
    return False
  return (
      (GetFileType(handle) & ~FILE_TYPE_REMOTE) == FILE_TYPE_CHAR and
       GetConsoleMode(handle, byref(DWORD())))


def win_get_unicode_stream(stream, excepted_fileno, output_handle, encoding):
  """Returns a unicode-compatible stream.

  This function will return a direct-Console writing object only if:
  - the file number is the expected console file number
  - the handle the expected file handle
  - the 'real' handle is in fact a handle to a console.
  """
  old_fileno = getattr(stream, 'fileno', lambda: None)()
  if old_fileno == excepted_fileno:
306
    # These types are available on linux but not Mac.
307
    # pylint: disable=no-name-in-module,F0401
308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350
    from ctypes import windll, WINFUNCTYPE
    from ctypes.wintypes import DWORD, HANDLE

    # <http://msdn.microsoft.com/en-us/library/ms683231.aspx>
    GetStdHandle = WINFUNCTYPE(HANDLE, DWORD)(('GetStdHandle', windll.kernel32))

    real_output_handle = GetStdHandle(DWORD(output_handle))
    if win_handle_is_a_console(real_output_handle):
      # It's a console.
      return WinUnicodeConsoleOutput(
          real_output_handle, old_fileno, stream.name, encoding)

  # It's something else. Create an auto-encoding stream.
  return WinUnicodeOutput(stream, old_fileno, encoding)


def fix_win_console(encoding):
  """Makes Unicode console output work independently of the current code page.

  This also fixes <http://bugs.python.org/issue1602>.
  Credit to Michael Kaplan
  <http://blogs.msdn.com/b/michkap/archive/2010/04/07/9989346.aspx> and
  TZOmegaTZIOY
  <http://stackoverflow.com/questions/878972/windows-cmd-encoding-change-causes-python-crash/1432462#1432462>.
  """
  if (isinstance(sys.stdout, WinUnicodeOutputBase) or
      isinstance(sys.stderr, WinUnicodeOutputBase)):
    return False

  try:
    # SetConsoleCP and SetConsoleOutputCP could be used to change the code page
    # but it's not really useful since the code here is using WriteConsoleW().
    # Also, changing the code page is 'permanent' to the console and needs to be
    # reverted manually.
    # In practice one needs to set the console font to a TTF font to be able to
    # see all the characters but it failed for me in practice. In any case, it
    # won't throw any exception when printing, which is the important part.
    # -11 and -12 are defined in stdio.h
    sys.stdout = win_get_unicode_stream(sys.stdout, 1, -11, encoding)
    sys.stderr = win_get_unicode_stream(sys.stderr, 2, -12, encoding)
    # TODO(maruel): Do sys.stdin with ReadConsoleW(). Albeit the limitation is
    # "It doesn't appear to be possible to read Unicode characters in UTF-8
    # mode" and this appears to be a limitation of cmd.exe.
351
  except Exception as e:
352 353 354 355 356 357 358
    complain('exception %r while fixing up sys.stdout and sys.stderr' % e)
  return True


def fix_encoding():
  """Fixes various encoding problems on all platforms.

359
  Should be called at the very beginning of the process.
360 361 362 363 364 365 366 367 368 369 370 371
  """
  ret = True
  if sys.platform == 'win32':
    ret &= fix_win_codec()

  ret &= fix_default_encoding()

  if sys.platform == 'win32':
    encoding = sys.getdefaultencoding()
    ret &= fix_win_sys_argv(encoding)
    ret &= fix_win_console(encoding)
  return ret