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