setuptools-0.6c16dev6.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         def path_append(value, addition):
 158             if value != "":
 159                 return value + u';' + addition
 160             else:
 161                 return addition
 162
 163         if addition.lower() in value.lower().split(u';'):
 164             print "Already done: %s." % (goal,)
 165         else:
 166             changed |= update(key_name_path, varname, path_append(value, addition), type, goal, what)
 167
 168         if change_allusers:
 169             # Also change any overriding environment entry for the current user.
 170             (user_value, user_type) = user_valueandtype or (u'', REG_SZ)
 171             split_value = user_value.lower().split(u';')
 172
 173             if not (addition.lower() in split_value or u'%'+varname.lower()+u'%' in split_value):
 174                 now_what = "the overriding user environment variable %s" % (varname,)
 175                 changed |= update((HKEY_CURRENT_USER, "HKEY_CURRENT_USER", USER_ENV),
 176                                   varname, path_append(user_value, addition), user_type,
 177                                   "add %s to %s" % (addition, now_what), now_what)
 178
 179         return changed
 180
 181
 182     def associate(ext, target, change_allusers):
 183         changed = False
 184         what = "the %s association for %s" % (change_allusers and "system" or "user", ext)
 185         goal = "associate the filetype %s with %s for %s" % (ext, target, change_allusers and "all users" or "the current user")
 186
 187         try:
 188             if change_allusers:
 189                 target_key = OpenKey(HKEY_LOCAL_MACHINE, "%s\\%s" % (SYSTEM_CLASSES, target), 0, KEY_QUERY_VALUE)
 190             else:
 191                 target_key = OpenKey(HKEY_CLASSES_ROOT, target, 0, KEY_QUERY_VALUE)
 192         except WindowsError, e:
 193             raise DistutilsSetupError("I was going to %s, but that won't work because the %s class does not exist in the registry, "
 194                                       "as far as I can tell.\n%r" % (goal, target, e))
 195         CloseKey(target_key)
 196
 197         system_key_name_path = (HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE", "%s\\%s" % (SYSTEM_CLASSES, ext))
 198         user_key_name_path   = (HKEY_CURRENT_USER,  "HKEY_CURRENT_USER",  "%s\\%s" % (USER_CLASSES,   ext))
 199
 200         system_valueandtype = open_and_query(system_classes, ext, "", "the system association for %s" % (ext,))
 201         user_valueandtype   = open_and_query(user_classes,   ext, "", "the user association for %s" % (ext,))
 202
 203         if change_allusers:
 204             (value, type) = system_valueandtype or (u'', REG_SZ)
 205             key_name_path = system_key_name_path
 206         else:
 207             (value, type) = user_valueandtype or system_valueandtype or (u'', REG_SZ)
 208             key_name_path = user_key_name_path
 209
 210         if value == target:
 211             print "Already done: %s." % (goal,)
 212         else:
 213             changed |= update(key_name_path, "", unicode(target), REG_SZ, goal, what)
 214
 215         if change_allusers:
 216             # Also change any overriding association for the current user.
 217             (user_value, user_type) = user_valueandtype or (u'', REG_SZ)
 218
 219             if user_value != target:
 220                 changed |= update(user_key_name_path, "", unicode(target), REG_SZ,
 221                                   "associate the filetype %s with %s for the current user " \
 222                                       "(because the system association is overridden)" % (ext, target),
 223                                   "the overriding user association for %s" % (ext,))
 224
 225         return changed
 226
 227
 228     def broadcast_settingchange(change_allusers):
 229         print "Broadcasting that the environment has changed, please wait..."
 230
 231         # <http://support.microsoft.com/kb/104011/en-us>
 232         # <http://msdn.microsoft.com/en-us/library/ms644952(VS.85).aspx>
 233         # LRESULT WINAPI SendMessageTimeoutW(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam,
 234         #                                    UINT fuFlags, UINT uTimeout, PDWORD_PTR lpdwResult);
 235
 236         try:
 237             from ctypes import WINFUNCTYPE, POINTER, windll, addressof, c_wchar_p
 238             from ctypes.wintypes import LONG, HWND, UINT, WPARAM, LPARAM, DWORD
 239
 240             SendMessageTimeout = WINFUNCTYPE(POINTER(LONG), HWND, UINT, WPARAM, LPARAM, UINT, UINT, POINTER(POINTER(DWORD))) \
 241                                      (("SendMessageTimeoutW", windll.user32))
 242             HWND_BROADCAST   = 0xFFFF
 243             WM_SETTINGCHANGE = 0x001A
 244             SMTO_ABORTIFHUNG = 0x0002
 245             SendMessageTimeout(HWND_BROADCAST, WM_SETTINGCHANGE, change_allusers and 1 or 0,
 246                                addressof(c_wchar_p(u"Environment")), SMTO_ABORTIFHUNG, 5000, None);
 247         except Exception, e:
 248             print "Warning: %r" % (e,)
 249
 250
 251     changed_assoc = associate(".pyscript", "Python.File", allusers)
 252
 253     changed_env = False
 254     try:
 255         changed_env |= add_to_environment("PATHEXT", ".pyscript", allusers)
 256         changed_env |= add_to_environment("PATHEXT", ".pyw",      allusers)
 257     finally:
 258         CloseKey(user_env)
 259         CloseKey(system_env)
 260
 261     if changed_assoc or changed_env:
 262         broadcast_settingchange(allusers)
 263
 264     if changed_env:
 265         # whether logout is needed seems to randomly differ between installations
 266         # of XP, but it is not needed in Vista or later.
 267         try:
 268             import platform, re
 269             need_logout = not re.search(r'^[6-9]|([1-9][0-9]+)\.', platform.version())
 270         except Exception, e:
 271             e  # hush pyflakes
 272             need_logout = True
 273
 274         if need_logout:
 275             print """
 276 ***********************************************************************
 277 Changes have been made to the persistent environment, but they may not
 278 take effect in this Windows session. Running installed Python scripts
 279 from a Command Prompt may only work after you have logged out and back
 280 in again, or rebooted.
 281 ***********************************************************************
 282 """
 283         else:
 284             print """
 285 ***********************************************************************
 286 Changes have been made to the persistent environment, but not in this
 287 Command Prompt. Running installed Python scripts will only work from
 288 new Command Prompts opened from now on.
 289 ***********************************************************************
 290 """