src/allmydata/windows/fixups.py

   1
   2 done = False
   3
   4 def initialize():
   5     global done
   6     import sys
   7     if sys.platform != "win32" or done:
   8         return True
   9     done = True
  10
  11     original_stderr = sys.stderr
  12     import codecs, re
  13     from ctypes import WINFUNCTYPE, windll, POINTER, byref, c_int
  14     from ctypes.wintypes import BOOL, HANDLE, DWORD, LPWSTR, LPCWSTR, LPVOID
  15     from allmydata.util import log
  16     from allmydata.util.encodingutil import canonical_encoding
  17
  18     # Work around <http://bugs.python.org/issue6058>.
  19     codecs.register(lambda name: name == 'cp65001' and codecs.lookup('utf-8') or None)
  20
  21     # Make Unicode console output work independently of the current code page.
  22     # This also fixes <http://bugs.python.org/issue1602>.
  23     # Credit to Michael Kaplan <http://blogs.msdn.com/b/michkap/archive/2010/04/07/9989346.aspx>
  24     # and TZOmegaTZIOY
  25     # <http://stackoverflow.com/questions/878972/windows-cmd-encoding-change-causes-python-crash/1432462#1432462>.
  26     try:
  27         # <http://msdn.microsoft.com/en-us/library/ms683231(VS.85).aspx>
  28         # HANDLE WINAPI GetStdHandle(DWORD nStdHandle);
  29         # returns INVALID_HANDLE_VALUE, NULL, or a valid handle
  30         #
  31         # <http://msdn.microsoft.com/en-us/library/aa364960(VS.85).aspx>
  32         # DWORD WINAPI GetFileType(DWORD hFile);
  33         #
  34         # <http://msdn.microsoft.com/en-us/library/ms683167(VS.85).aspx>
  35         # BOOL WINAPI GetConsoleMode(HANDLE hConsole, LPDWORD lpMode);
  36
  37         GetStdHandle = WINFUNCTYPE(HANDLE, DWORD)(("GetStdHandle", windll.kernel32))
  38         STD_OUTPUT_HANDLE = DWORD(-11)
  39         STD_ERROR_HANDLE  = DWORD(-12)
  40         GetFileType = WINFUNCTYPE(DWORD, DWORD)(("GetFileType", windll.kernel32))
  41         FILE_TYPE_CHAR   = 0x0002
  42         FILE_TYPE_REMOTE = 0x8000
  43         GetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, POINTER(DWORD))(("GetConsoleMode", windll.kernel32))
  44         INVALID_HANDLE_VALUE = DWORD(-1).value
  45
  46         def not_a_console(handle):
  47             if handle == INVALID_HANDLE_VALUE or handle is None:
  48                 return True
  49             return ((GetFileType(handle) & ~FILE_TYPE_REMOTE) != FILE_TYPE_CHAR
  50                     or GetConsoleMode(handle, byref(DWORD())) == 0)
  51
  52         old_stdout_fileno = None
  53         old_stderr_fileno = None
  54         if hasattr(sys.stdout, 'fileno'):
  55             old_stdout_fileno = sys.stdout.fileno()
  56         if hasattr(sys.stderr, 'fileno'):
  57             old_stderr_fileno = sys.stderr.fileno()
  58
  59         STDOUT_FILENO = 1
  60         STDERR_FILENO = 2
  61         real_stdout = (old_stdout_fileno == STDOUT_FILENO)
  62         real_stderr = (old_stderr_fileno == STDERR_FILENO)
  63
  64         if real_stdout:
  65             hStdout = GetStdHandle(STD_OUTPUT_HANDLE)
  66             if not_a_console(hStdout):
  67                 real_stdout = False
  68
  69         if real_stderr:
  70             hStderr = GetStdHandle(STD_ERROR_HANDLE)
  71             if not_a_console(hStderr):
  72                 real_stderr = False
  73
  74         if real_stdout or real_stderr:
  75             # BOOL WINAPI WriteConsoleW(HANDLE hOutput, LPWSTR lpBuffer, DWORD nChars,
  76             #                           LPDWORD lpCharsWritten, LPVOID lpReserved);
  77
  78             WriteConsoleW = WINFUNCTYPE(BOOL, HANDLE, LPWSTR, DWORD, POINTER(DWORD), LPVOID) \
  79                                 (("WriteConsoleW", windll.kernel32))
  80
  81             # If any exception occurs in this code, we'll probably try to print it on stderr,
  82             # which makes for frustrating debugging if stderr is directed to this code.
  83             # So be paranoid about catching errors and reporting them to original_stderr,
  84             # so that we can at least see them.
  85
  86             class UnicodeOutput:
  87                 def __init__(self, hConsole, stream, fileno, name):
  88                     self._hConsole = hConsole
  89                     self._stream = stream
  90                     self._fileno = fileno
  91                     self.closed = False
  92                     self.softspace = False
  93                     self.mode = 'w'
  94                     self.encoding = 'utf-8'
  95                     self.name = name
  96                     if hasattr(stream, 'encoding') and canonical_encoding(stream.encoding) != 'utf-8':
  97                         log.msg("%s (%r) had encoding %r, but we're going to write UTF-8 to it" %
  98                                 (name, stream, stream.encoding), level=log.CURIOUS)
  99                     self.flush()
 100
 101                 def isatty(self):
 102                     return False
 103                 def close(self):
 104                     # don't really close the handle, that would only cause problems
 105                     self.closed = True
 106                 def fileno(self):
 107                     return self._fileno
 108                 def flush(self):
 109                     if self._hConsole is None:
 110                         try:
 111                             self._stream.flush()
 112                         except Exception, e:
 113                             print >>original_stderr, repr(e)
 114                             raise
 115
 116                 def write(self, text):
 117                     try:
 118                         if self._hConsole is None:
 119                             if isinstance(text, unicode):
 120                                 text = text.encode('utf-8')
 121                             self._stream.write(text)
 122                         else:
 123                             if not isinstance(text, unicode):
 124                                 text = str(text).decode('utf-8')
 125                             remaining = len(text)
 126                             while remaining > 0:
 127                                 n = DWORD(0)
 128                                 retval = WriteConsoleW(self._hConsole, text, remaining, byref(n), None)
 129                                 if retval == 0 or n.value == 0:
 130                                     raise IOError("could not write to %s [WriteConsoleW returned %r, n.value = %r]"
 131                                                   % (self.name, retval, n.value))
 132                                 remaining -= n.value
 133                                 if remaining == 0: break
 134                                 text = text[n.value:]
 135                     except Exception, e:
 136                         print >>original_stderr, repr(e)
 137                         raise
 138
 139                 def writelines(self, lines):
 140                     try:
 141                         for line in lines:
 142                             self.write(line)
 143                     except Exception, e:
 144                         print >>original_stderr, repr(e)
 145                         raise
 146
 147             if real_stdout:
 148                 sys.stdout = UnicodeOutput(hStdout, None, STDOUT_FILENO, '<Unicode console stdout>')
 149             else:
 150                 sys.stdout = UnicodeOutput(None, sys.stdout, old_stdout_fileno, '<Unicode redirected stdout>')
 151
 152             if real_stderr:
 153                 sys.stderr = UnicodeOutput(hStderr, None, STDERR_FILENO, '<Unicode console stderr>')
 154             else:
 155                 sys.stderr = UnicodeOutput(None, sys.stderr, old_stderr_fileno, '<Unicode redirected stdout>')
 156     except Exception, e:
 157         print >>original_stderr, "exception %r while fixing up sys.stdout and sys.stderr" % (e,)
 158         log.msg("exception %r while fixing up sys.stdout and sys.stderr" % (e,), log.WEIRD)
 159
 160     # Unmangle command-line arguments.
 161     GetCommandLineW = WINFUNCTYPE(LPWSTR)(("GetCommandLineW", windll.kernel32))
 162     CommandLineToArgvW = WINFUNCTYPE(POINTER(LPWSTR), LPCWSTR, POINTER(c_int)) \
 163                             (("CommandLineToArgvW", windll.shell32))
 164
 165     argc = c_int(0)
 166     argv_unicode = CommandLineToArgvW(GetCommandLineW(), byref(argc))
 167
 168     def unmangle(s):
 169         return re.sub(ur'\x7f[0-9a-fA-F]*\;', lambda m: unichr(int(m.group(0)[1:-1], 16)), s)
 170
 171     try:
 172         sys.argv = [unmangle(argv_unicode[i]).encode('utf-8') for i in xrange(1, argc.value)]
 173     except Exception, e:
 174         print >>sys.stderr, "%s:  could not unmangle Unicode arguments" % (sys.argv[0],)
 175         print >>sys.stderr, [argv_unicode[i] for i in xrange(1, argc.value)]
 176         raise
 177
 178     if sys.argv[0].endswith('.pyscript'):
 179         sys.argv[0] = sys.argv[0][:-9]