zfec/setuptools-0.6c16dev3.egg/setuptools/command/scriptsetup.py

   1 from distutils.errors import DistutilsSetupError
   2 from setuptools import Command
   3 import sys
   4
   5 class scriptsetup(Command):
   6     action = (sys.platform == "win32"
   7                    and "set up .pyscript association and PATHEXT variable to run scripts"
   8                    or  "this does nothing on non-Windows platforms")
   9
  10     user_options = [
  11         ('allusers', 'a',
  12          'make changes for all users of this Windows installation (requires Administrator privileges)'),
  13     ]
  14     boolean_options = ['allusers']
  15
  16     def initialize_options(self):
  17         self.allusers = False
  18
  19     def finalize_options(self):
  20         pass
  21
  22     def run(self):
  23         if sys.platform != "win32":
  24             print "\n'scriptsetup' isn't needed on non-Windows platforms."
  25         else:
  26             do_scriptsetup(self.allusers)
  27
  28
  29 def do_scriptsetup(allusers=False):
  30     print "\nSetting up environment to run scripts for %s..." % (allusers and "all users" or "the current user")
  31
  32     from _winreg import HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, HKEY_CLASSES_ROOT, \
  33         REG_SZ, REG_EXPAND_SZ, KEY_QUERY_VALUE, KEY_SET_VALUE, \
  34         OpenKey, CreateKey, QueryValueEx, SetValueEx, FlushKey, CloseKey
  35
  36     USER_ENV = "Environment"
  37     try:
  38         user_env = OpenKey(HKEY_CURRENT_USER, USER_ENV, 0, KEY_QUERY_VALUE)
  39     except WindowsError, e:
  40         raise DistutilsSetupError("I could not read the user environment from the registry.\n%r" % (e,))
  41
  42     SYSTEM_ENV = "SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment"
  43     try:
  44         system_env = OpenKey(HKEY_LOCAL_MACHINE, SYSTEM_ENV, 0, KEY_QUERY_VALUE)
  45     except WindowsError, e:
  46         raise DistutilsSetupError("I could not read the system environment from the registry.\n%r" % (e,))
  47
  48
  49     # HKEY_CLASSES_ROOT is a merged view that would only confuse us.
  50     # <http://technet.microsoft.com/en-us/library/cc739822(WS.10).aspx>
  51
  52     USER_CLASSES = "SOFTWARE\\Classes"
  53     try:
  54         user_classes = OpenKey(HKEY_CURRENT_USER, USER_CLASSES, 0, KEY_QUERY_VALUE)
  55     except WindowsError, e:
  56         raise DistutilsSetupError("I could not read the user filetype associations from the registry.\n%r" % (e,))
  57
  58     SYSTEM_CLASSES = "SOFTWARE\\Classes"
  59     try:
  60         system_classes = OpenKey(HKEY_LOCAL_MACHINE, SYSTEM_CLASSES, 0, KEY_QUERY_VALUE)
  61     except WindowsError, e:
  62         raise DistutilsSetupError("I could not read the system filetype associations from the registry.\n%r" % (e,))
  63
  64
  65     def query(key, subkey, what):
  66         try:
  67             (value, type) = QueryValueEx(key, subkey)
  68         except WindowsError, e:
  69             if e.winerror == 2:  # not found
  70                 return None
  71             raise DistutilsSetupError("I could not read %s from the registry.\n%r" % (what, e))
  72
  73         # It does not matter that we don't expand environment strings, in fact it's better not to.
  74
  75         if type != REG_SZ and type != REG_EXPAND_SZ:
  76             raise DistutilsSetupError("I expected the registry entry for %s to have a string type (REG_SZ or REG_EXPAND_SZ), "
  77                                       "and was flummoxed by it having type code %r." % (what, type))
  78         return (value, type)
  79
  80
  81     def open_and_query(key, path, subkey, what):
  82         try:
  83             read_key = OpenKey(key, path, 0, KEY_QUERY_VALUE)
  84         except WindowsError, e:
  85             if e.winerror == 2:  # not found
  86                 return None
  87             raise DistutilsSetupError("I could not read %s from the registry because I could not open "
  88                                       "the parent key.\n%r" % (what, e))
  89
  90         try:
  91             return query(read_key, subkey, what)
  92         finally:
  93             CloseKey(read_key)
  94
  95
  96     def update(key_name_path, subkey, desired_value, desired_type, goal, what):
  97         (key, name, path) = key_name_path
  98
  99         (old_value, old_type) = open_and_query(key, path, subkey, what) or (None, None)
 100         if (old_value, old_type) == (desired_value, desired_type):
 101             print "Already done: %s." % (goal,)
 102             return False
 103
 104         try:
 105             update_key = OpenKey(key, path, 0, KEY_SET_VALUE|KEY_QUERY_VALUE)
 106         except WindowsError, e:
 107             if e.winerror != 2:
 108                 raise DistutilsSetupError("I tried to %s, but was not successful because I could not open "
 109                                           "the registry key %s\\%s for writing.\n%r"
 110                                           % (goal, name, path, e))
 111             try:
 112                 update_key = CreateKey(key, path)
 113             except WindowsError, e:
 114                 raise DistutilsSetupError("I tried to %s, but was not successful because the registry key %s\\%s "
 115                                           "did not exist, and I was unable to create it.\n%r"
 116                                           % (goal, name, path, e))
 117
 118         (new_value, new_type) = (None, None)
 119         try:
 120             SetValueEx(update_key, subkey, 0, desired_type, desired_value)
 121         except WindowsError, e:
 122             raise DistutilsSetupError("I tried to %s, but was not able to set the subkey %r under %s\\%s to be %r.\n%r"
 123                                       % (goal, subkey, name, path, desired_value))
 124         else:
 125             (new_value, new_type) = query(update_key, subkey, what) or (None, None)
 126         finally:
 127             FlushKey(update_key)
 128             CloseKey(update_key)
 129
 130         if (new_value, new_type) != (desired_value, desired_type):
 131             raise DistutilsSetupError("I tried to %s by setting the subkey %r under %s\\%s to be %r, "
 132                                       "and the call to SetValueEx succeeded, but the value ended up as "
 133                                       "%r instead (it was previously %r). Maybe the update was unexpectedly virtualized?"
 134                                       % (goal, subkey, name, path, desired_value, new_value, old_value))
 135
 136         print "Done: %s." % (goal,)
 137         return True
 138
 139
 140     # Maintenance hazard: 'add_to_environment' and 'associate' use very similar, but not identical logic.
 141
 142     def add_to_environment(varname, addition, change_allusers):
 143         changed = False
 144         what = "the %s environment variable %s" % (change_allusers and "system" or "user", varname)
 145         goal = "add %s to %s" % (addition, what)
 146
 147         system_valueandtype = query(system_env, varname, "the system environment variable %s" % (varname,))
 148         user_valueandtype   = query(user_env,   varname, "the user environment variable %s" % (varname,))
 149
 150         if change_allusers:
 151             (value, type) = system_valueandtype or (u'', REG_SZ)
 152             key_name_path = (HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE", SYSTEM_ENV)
 153         else:
 154             (value, type) = user_valueandtype or system_valueandtype or (u'', REG_SZ)
 155             key_name_path = (HKEY_CURRENT_USER, "HKEY_CURRENT_USER", USER_ENV)
 156
 157         if addition.lower() in value.lower().split(u';'):
 158             print "Already done: %s." % (goal,)
 159         else:
 160             changed |= update(key_name_path, varname, value + u';' + addition, type, goal, what)
 161
 162         if change_allusers:
 163             # Also change any overriding environment entry for the current user.
 164             (user_value, user_type) = user_valueandtype or (u'', REG_SZ)
 165             split_value = user_value.lower().split(u';')
 166
 167             if not (addition.lower() in split_value or u'%'+varname.lower()+u'%' in split_value):
 168                 now_what = "the overriding user environment variable %s" % (varname,)
 169                 changed |= update((HKEY_CURRENT_USER, "HKEY_CURRENT_USER", USER_ENV),
 170                                   varname, user_value + u';' + addition, user_type,
 171                                   "add %s to %s" % (addition, now_what), now_what)
 172
 173         return changed
 174
 175
 176     def associate(ext, target, change_allusers):
 177         changed = False
 178         what = "the %s association for %s" % (change_allusers and "system" or "user", ext)
 179         goal = "associate the filetype %s with %s for %s" % (ext, target, change_allusers and "all users" or "the current user")
 180
 181         try:
 182             if change_allusers:
 183                 target_key = OpenKey(HKEY_LOCAL_MACHINE, "%s\\%s" % (SYSTEM_CLASSES, target), 0, KEY_QUERY_VALUE)
 184             else:
 185                 target_key = OpenKey(HKEY_CLASSES_ROOT, target, 0, KEY_QUERY_VALUE)
 186         except WindowsError, e:
 187             raise DistutilsSetupError("I was going to %s, but that won't work because the %s class does not exist in the registry, "
 188                                       "as far as I can tell.\n%r" % (goal, target, e))
 189         CloseKey(target_key)
 190
 191         system_key_name_path = (HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE", "%s\\%s" % (SYSTEM_CLASSES, ext))
 192         user_key_name_path   = (HKEY_CURRENT_USER,  "HKEY_CURRENT_USER",  "%s\\%s" % (USER_CLASSES,   ext))
 193
 194         system_valueandtype = open_and_query(system_classes, ext, "", "the system association for %s" % (ext,))
 195         user_valueandtype   = open_and_query(user_classes,   ext, "", "the user association for %s" % (ext,))
 196
 197         if change_allusers:
 198             (value, type) = system_valueandtype or (u'', REG_SZ)
 199             key_name_path = system_key_name_path
 200         else:
 201             (value, type) = user_valueandtype or system_valueandtype or (u'', REG_SZ)
 202             key_name_path = user_key_name_path
 203
 204         if value == target:
 205             print "Already done: %s." % (goal,)
 206         else:
 207             changed |= update(key_name_path, "", unicode(target), REG_SZ, goal, what)
 208
 209         if change_allusers:
 210             # Also change any overriding association for the current user.
 211             (user_value, user_type) = user_valueandtype or (u'', REG_SZ)
 212
 213             if user_value != target:
 214                 changed |= update(user_key_name_path, "", unicode(target), REG_SZ,
 215                                   "associate the filetype %s with %s for the current user " \
 216                                       "(because the system association is overridden)" % (ext, target),
 217                                   "the overriding user association for %s" % (ext,))
 218
 219         return changed
 220
 221
 222     def broadcast_settingchange(change_allusers):
 223         print "Broadcasting that the environment has changed, please wait..."
 224
 225         # <http://support.microsoft.com/kb/104011/en-us>
 226         # <http://msdn.microsoft.com/en-us/library/ms644952(VS.85).aspx>
 227         # LRESULT WINAPI SendMessageTimeoutW(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam,
 228         #                                    UINT fuFlags, UINT uTimeout, PDWORD_PTR lpdwResult);
 229
 230         try:
 231             from ctypes import WINFUNCTYPE, POINTER, windll, addressof, c_wchar_p
 232             from ctypes.wintypes import LONG, HWND, UINT, WPARAM, LPARAM, DWORD
 233
 234             SendMessageTimeout = WINFUNCTYPE(POINTER(LONG), HWND, UINT, WPARAM, LPARAM, UINT, UINT, POINTER(POINTER(DWORD))) \
 235                                      (("SendMessageTimeoutW", windll.user32))
 236             HWND_BROADCAST   = 0xFFFF
 237             WM_SETTINGCHANGE = 0x001A
 238             SMTO_ABORTIFHUNG = 0x0002
 239             SendMessageTimeout(HWND_BROADCAST, WM_SETTINGCHANGE, change_allusers and 1 or 0,
 240                                addressof(c_wchar_p(u"Environment")), SMTO_ABORTIFHUNG, 5000, None);
 241         except Exception, e:
 242             print "Warning: %r" % (e,)
 243
 244
 245     changed_assoc = associate(".pyscript", "Python.File", allusers)
 246
 247     changed_env = False
 248     try:
 249         changed_env |= add_to_environment("PATHEXT", ".pyscript", allusers)
 250         changed_env |= add_to_environment("PATHEXT", ".pyw",      allusers)
 251     finally:
 252         CloseKey(user_env)
 253         CloseKey(system_env)
 254
 255     if changed_assoc or changed_env:
 256         broadcast_settingchange(allusers)
 257
 258     if changed_env:
 259         # whether logout is needed seems to randomly differ between installations
 260         # of XP, but it is not needed in Vista or later.
 261         try:
 262             import platform, re
 263             need_logout = not re.search(r'^[6-9]|([1-9][0-9]+)\.', platform.version())
 264         except Exception, e:
 265             e  # hush pyflakes
 266             need_logout = True
 267
 268         if need_logout:
 269             print """
 270 ***********************************************************************
 271 Changes have been made to the persistent environment, but they may not
 272 take effect in this Windows session. Running installed Python scripts
 273 from a Command Prompt may only work after you have logged out and back
 274 in again, or rebooted.
 275 ***********************************************************************
 276 """
 277         else:
 278             print """
 279 ***********************************************************************
 280 Changes have been made to the persistent environment, but not in this
 281 Command Prompt. Running installed Python scripts will only work from
 282 new Command Prompts opened from now on.
 283 ***********************************************************************
 284 """