eventlet/eventlet/patcher.py

   1 import imp
   2 import sys
   3
   4 from eventlet.support import six
   5
   6
   7 __all__ = ['inject', 'import_patched', 'monkey_patch', 'is_monkey_patched']
   8
   9 __exclude = set(('__builtins__', '__file__', '__name__'))
  10
  11
  12 class SysModulesSaver(object):
  13     """Class that captures some subset of the current state of
  14     sys.modules.  Pass in an iterator of module names to the
  15     constructor."""
  16
  17     def __init__(self, module_names=()):
  18         self._saved = {}
  19         imp.acquire_lock()
  20         self.save(*module_names)
  21
  22     def save(self, *module_names):
  23         """Saves the named modules to the object."""
  24         for modname in module_names:
  25             self._saved[modname] = sys.modules.get(modname, None)
  26
  27     def restore(self):
  28         """Restores the modules that the saver knows about into
  29         sys.modules.
  30         """
  31         try:
  32             for modname, mod in six.iteritems(self._saved):
  33                 if mod is not None:
  34                     sys.modules[modname] = mod
  35                 else:
  36                     try:
  37                         del sys.modules[modname]
  38                     except KeyError:
  39                         pass
  40         finally:
  41             imp.release_lock()
  42
  43
  44 def inject(module_name, new_globals, *additional_modules):
  45     """Base method for "injecting" greened modules into an imported module.  It
  46     imports the module specified in *module_name*, arranging things so
  47     that the already-imported modules in *additional_modules* are used when
  48     *module_name* makes its imports.
  49
  50     *new_globals* is either None or a globals dictionary that gets populated
  51     with the contents of the *module_name* module.  This is useful when creating
  52     a "green" version of some other module.
  53
  54     *additional_modules* should be a collection of two-element tuples, of the
  55     form (<name>, <module>).  If it's not specified, a default selection of
  56     name/module pairs is used, which should cover all use cases but may be
  57     slower because there are inevitably redundant or unnecessary imports.
  58     """
  59     patched_name = '__patched_module_' + module_name
  60     if patched_name in sys.modules:
  61         # returning already-patched module so as not to destroy existing
  62         # references to patched modules
  63         return sys.modules[patched_name]
  64
  65     if not additional_modules:
  66         # supply some defaults
  67         additional_modules = (
  68             _green_os_modules() +
  69             _green_select_modules() +
  70             _green_socket_modules() +
  71             _green_thread_modules() +
  72             _green_time_modules())
  73             # _green_MySQLdb()) # enable this after a short baking-in period
  74
  75     # after this we are gonna screw with sys.modules, so capture the
  76     # state of all the modules we're going to mess with, and lock
  77     saver = SysModulesSaver([name for name, m in additional_modules])
  78     saver.save(module_name)
  79
  80     # Cover the target modules so that when you import the module it
  81     # sees only the patched versions
  82     for name, mod in additional_modules:
  83         sys.modules[name] = mod
  84
  85     # Remove the old module from sys.modules and reimport it while
  86     # the specified modules are in place
  87     sys.modules.pop(module_name, None)
  88     try:
  89         module = __import__(module_name, {}, {}, module_name.split('.')[:-1])
  90
  91         if new_globals is not None:
  92             # Update the given globals dictionary with everything from this new module
  93             for name in dir(module):
  94                 if name not in __exclude:
  95                     new_globals[name] = getattr(module, name)
  96
  97         # Keep a reference to the new module to prevent it from dying
  98         sys.modules[patched_name] = module
  99     finally:
 100         saver.restore()  # Put the original modules back
 101
 102     return module
 103
 104
 105 def import_patched(module_name, *additional_modules, **kw_additional_modules):
 106     """Imports a module in a way that ensures that the module uses "green"
 107     versions of the standard library modules, so that everything works
 108     nonblockingly.
 109
 110     The only required argument is the name of the module to be imported.
 111     """
 112     return inject(
 113         module_name,
 114         None,
 115         *additional_modules + tuple(kw_additional_modules.items()))
 116
 117
 118 def patch_function(func, *additional_modules):
 119     """Decorator that returns a version of the function that patches
 120     some modules for the duration of the function call.  This is
 121     deeply gross and should only be used for functions that import
 122     network libraries within their function bodies that there is no
 123     way of getting around."""
 124     if not additional_modules:
 125         # supply some defaults
 126         additional_modules = (
 127             _green_os_modules() +
 128             _green_select_modules() +
 129             _green_socket_modules() +
 130             _green_thread_modules() +
 131             _green_time_modules())
 132
 133     def patched(*args, **kw):
 134         saver = SysModulesSaver()
 135         for name, mod in additional_modules:
 136             saver.save(name)
 137             sys.modules[name] = mod
 138         try:
 139             return func(*args, **kw)
 140         finally:
 141             saver.restore()
 142     return patched
 143
 144
 145 def _original_patch_function(func, *module_names):
 146     """Kind of the contrapositive of patch_function: decorates a
 147     function such that when it's called, sys.modules is populated only
 148     with the unpatched versions of the specified modules.  Unlike
 149     patch_function, only the names of the modules need be supplied,
 150     and there are no defaults.  This is a gross hack; tell your kids not
 151     to import inside function bodies!"""
 152     def patched(*args, **kw):
 153         saver = SysModulesSaver(module_names)
 154         for name in module_names:
 155             sys.modules[name] = original(name)
 156         try:
 157             return func(*args, **kw)
 158         finally:
 159             saver.restore()
 160     return patched
 161
 162
 163 def original(modname):
 164     """ This returns an unpatched version of a module; this is useful for
 165     Eventlet itself (i.e. tpool)."""
 166     # note that it's not necessary to temporarily install unpatched
 167     # versions of all patchable modules during the import of the
 168     # module; this is because none of them import each other, except
 169     # for threading which imports thread
 170     original_name = '__original_module_' + modname
 171     if original_name in sys.modules:
 172         return sys.modules.get(original_name)
 173
 174     # re-import the "pure" module and store it in the global _originals
 175     # dict; be sure to restore whatever module had that name already
 176     saver = SysModulesSaver((modname,))
 177     sys.modules.pop(modname, None)
 178     # some rudimentary dependency checking -- fortunately the modules
 179     # we're working on don't have many dependencies so we can just do
 180     # some special-casing here
 181     if six.PY2:
 182         deps = {'threading': 'thread', 'Queue': 'threading'}
 183     if six.PY3:
 184         deps = {'threading': '_thread', 'queue': 'threading'}
 185     if modname in deps:
 186         dependency = deps[modname]
 187         saver.save(dependency)
 188         sys.modules[dependency] = original(dependency)
 189     try:
 190         real_mod = __import__(modname, {}, {}, modname.split('.')[:-1])
 191         if modname in ('Queue', 'queue') and not hasattr(real_mod, '_threading'):
 192             # tricky hack: Queue's constructor in <2.7 imports
 193             # threading on every instantiation; therefore we wrap
 194             # it so that it always gets the original threading
 195             real_mod.Queue.__init__ = _original_patch_function(
 196                 real_mod.Queue.__init__,
 197                 'threading')
 198         # save a reference to the unpatched module so it doesn't get lost
 199         sys.modules[original_name] = real_mod
 200     finally:
 201         saver.restore()
 202
 203     return sys.modules[original_name]
 204
 205 already_patched = {}
 206
 207
 208 def monkey_patch(**on):
 209     """Globally patches certain system modules to be greenthread-friendly.
 210
 211     The keyword arguments afford some control over which modules are patched.
 212     If no keyword arguments are supplied, all possible modules are patched.
 213     If keywords are set to True, only the specified modules are patched.  E.g.,
 214     ``monkey_patch(socket=True, select=True)`` patches only the select and
 215     socket modules.  Most arguments patch the single module of the same name
 216     (os, time, select).  The exceptions are socket, which also patches the ssl
 217     module if present; and thread, which patches thread, threading, and Queue.
 218
 219     It's safe to call monkey_patch multiple times.
 220     """
 221     accepted_args = set(('os', 'select', 'socket',
 222                          'thread', 'time', 'psycopg', 'MySQLdb', '__builtin__'))
 223     default_on = on.pop("all", None)
 224     for k in six.iterkeys(on):
 225         if k not in accepted_args:
 226             raise TypeError("monkey_patch() got an unexpected "
 227                             "keyword argument %r" % k)
 228     if default_on is None:
 229         default_on = not (True in on.values())
 230     for modname in accepted_args:
 231         if modname == 'MySQLdb':
 232             # MySQLdb is only on when explicitly patched for the moment
 233             on.setdefault(modname, False)
 234         if modname == '__builtin__':
 235             on.setdefault(modname, False)
 236         on.setdefault(modname, default_on)
 237
 238     modules_to_patch = []
 239     if on['os'] and not already_patched.get('os'):
 240         modules_to_patch += _green_os_modules()
 241         already_patched['os'] = True
 242     if on['select'] and not already_patched.get('select'):
 243         modules_to_patch += _green_select_modules()
 244         already_patched['select'] = True
 245     if on['socket'] and not already_patched.get('socket'):
 246         modules_to_patch += _green_socket_modules()
 247         already_patched['socket'] = True
 248     if on['thread'] and not already_patched.get('thread'):
 249         modules_to_patch += _green_thread_modules()
 250         already_patched['thread'] = True
 251     if on['time'] and not already_patched.get('time'):
 252         modules_to_patch += _green_time_modules()
 253         already_patched['time'] = True
 254     if on.get('MySQLdb') and not already_patched.get('MySQLdb'):
 255         modules_to_patch += _green_MySQLdb()
 256         already_patched['MySQLdb'] = True
 257     if on.get('__builtin__') and not already_patched.get('__builtin__'):
 258         modules_to_patch += _green_builtins()
 259         already_patched['__builtin__'] = True
 260     if on['psycopg'] and not already_patched.get('psycopg'):
 261         try:
 262             from eventlet.support import psycopg2_patcher
 263             psycopg2_patcher.make_psycopg_green()
 264             already_patched['psycopg'] = True
 265         except ImportError:
 266             # note that if we get an importerror from trying to
 267             # monkeypatch psycopg, we will continually retry it
 268             # whenever monkey_patch is called; this should not be a
 269             # performance problem but it allows is_monkey_patched to
 270             # tell us whether or not we succeeded
 271             pass
 272
 273     imp.acquire_lock()
 274     try:
 275         for name, mod in modules_to_patch:
 276             orig_mod = sys.modules.get(name)
 277             if orig_mod is None:
 278                 orig_mod = __import__(name)
 279             for attr_name in mod.__patched__:
 280                 patched_attr = getattr(mod, attr_name, None)
 281                 if patched_attr is not None:
 282                     setattr(orig_mod, attr_name, patched_attr)
 283     finally:
 284         imp.release_lock()
 285
 286
 287 def is_monkey_patched(module):
 288     """Returns True if the given module is monkeypatched currently, False if
 289     not.  *module* can be either the module itself or its name.
 290
 291     Based entirely off the name of the module, so if you import a
 292     module some other way than with the import keyword (including
 293     import_patched), this might not be correct about that particular
 294     module."""
 295     return module in already_patched or \
 296         getattr(module, '__name__', None) in already_patched
 297
 298
 299 def _green_os_modules():
 300     from eventlet.green import os
 301     return [('os', os)]
 302
 303
 304 def _green_select_modules():
 305     from eventlet.green import select
 306     return [('select', select)]
 307
 308
 309 def _green_socket_modules():
 310     from eventlet.green import socket
 311     try:
 312         from eventlet.green import ssl
 313         return [('socket', socket), ('ssl', ssl)]
 314     except ImportError:
 315         return [('socket', socket)]
 316
 317
 318 def _green_thread_modules():
 319     from eventlet.green import Queue
 320     from eventlet.green import thread
 321     from eventlet.green import threading
 322     if six.PY2:
 323         return [('Queue', Queue), ('thread', thread), ('threading', threading)]
 324     if six.PY3:
 325         return [('queue', Queue), ('_thread', thread), ('threading', threading)]
 326
 327
 328 def _green_time_modules():
 329     from eventlet.green import time
 330     return [('time', time)]
 331
 332
 333 def _green_MySQLdb():
 334     try:
 335         from eventlet.green import MySQLdb
 336         return [('MySQLdb', MySQLdb)]
 337     except ImportError:
 338         return []
 339
 340
 341 def _green_builtins():
 342     try:
 343         from eventlet.green import builtin
 344         return [('__builtin__', builtin)]
 345     except ImportError:
 346         return []
 347
 348
 349 def slurp_properties(source, destination, ignore=[], srckeys=None):
 350     """Copy properties from *source* (assumed to be a module) to
 351     *destination* (assumed to be a dict).
 352
 353     *ignore* lists properties that should not be thusly copied.
 354     *srckeys* is a list of keys to copy, if the source's __all__ is
 355     untrustworthy.
 356     """
 357     if srckeys is None:
 358         srckeys = source.__all__
 359     destination.update(dict([
 360         (name, getattr(source, name))
 361         for name in srckeys
 362         if not (name.startswith('__') or name in ignore)
 363     ]))
 364
 365
 366 if __name__ == "__main__":
 367     sys.argv.pop(0)
 368     monkey_patch()
 369     with open(sys.argv[0]) as f:
 370         code = compile(f.read(), sys.argv[0], 'exec')
 371         exec(code)