python-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',
 227                          'builtins'))
 228     # To make sure only one of them is passed here
 229     assert not ('__builtin__' in on and 'builtins' in on)
 230     try:
 231         b = on.pop('__builtin__')
 232     except KeyError:
 233         pass
 234     else:
 235         on['builtins'] = b
 236
 237     default_on = on.pop("all", None)
 238
 239     for k in six.iterkeys(on):
 240         if k not in accepted_args:
 241             raise TypeError("monkey_patch() got an unexpected "
 242                             "keyword argument %r" % k)
 243     if default_on is None:
 244         default_on = not (True in on.values())
 245     for modname in accepted_args:
 246         if modname == 'MySQLdb':
 247             # MySQLdb is only on when explicitly patched for the moment
 248             on.setdefault(modname, False)
 249         if modname == 'builtins':
 250             on.setdefault(modname, False)
 251         on.setdefault(modname, default_on)
 252
 253     modules_to_patch = []
 254     for name, modules_function in [
 255         ('os', _green_os_modules),
 256         ('select', _green_select_modules),
 257         ('socket', _green_socket_modules),
 258         ('thread', _green_thread_modules),
 259         ('time', _green_time_modules),
 260         ('MySQLdb', _green_MySQLdb),
 261         ('builtins', _green_builtins),
 262     ]:
 263         if on[name] and not already_patched.get(name):
 264             modules_to_patch += modules_function()
 265             already_patched[name] = True
 266
 267     if on['psycopg'] and not already_patched.get('psycopg'):
 268         try:
 269             from eventlet.support import psycopg2_patcher
 270             psycopg2_patcher.make_psycopg_green()
 271             already_patched['psycopg'] = True
 272         except ImportError:
 273             # note that if we get an importerror from trying to
 274             # monkeypatch psycopg, we will continually retry it
 275             # whenever monkey_patch is called; this should not be a
 276             # performance problem but it allows is_monkey_patched to
 277             # tell us whether or not we succeeded
 278             pass
 279
 280     imp.acquire_lock()
 281     try:
 282         for name, mod in modules_to_patch:
 283             orig_mod = sys.modules.get(name)
 284             if orig_mod is None:
 285                 orig_mod = __import__(name)
 286             for attr_name in mod.__patched__:
 287                 patched_attr = getattr(mod, attr_name, None)
 288                 if patched_attr is not None:
 289                     setattr(orig_mod, attr_name, patched_attr)
 290             deleted = getattr(mod, '__deleted__', [])
 291             for attr_name in deleted:
 292                 if hasattr(orig_mod, attr_name):
 293                     delattr(orig_mod, attr_name)
 294     finally:
 295         imp.release_lock()
 296
 297     if sys.version_info >= (3, 3):
 298         import importlib._bootstrap
 299         thread = original('_thread')
 300         # importlib must use real thread locks, not eventlet.Semaphore
 301         importlib._bootstrap._thread = thread
 302
 303         # Issue #185: Since Python 3.3, threading.RLock is implemented in C and
 304         # so call a C function to get the thread identifier, instead of calling
 305         # threading.get_ident(). Force the Python implementation of RLock which
 306         # calls threading.get_ident() and so is compatible with eventlet.
 307         import threading
 308         threading.RLock = threading._PyRLock
 309
 310
 311 def is_monkey_patched(module):
 312     """Returns True if the given module is monkeypatched currently, False if
 313     not.  *module* can be either the module itself or its name.
 314
 315     Based entirely off the name of the module, so if you import a
 316     module some other way than with the import keyword (including
 317     import_patched), this might not be correct about that particular
 318     module."""
 319     return module in already_patched or \
 320         getattr(module, '__name__', None) in already_patched
 321
 322
 323 def _green_os_modules():
 324     from eventlet.green import os
 325     return [('os', os)]
 326
 327
 328 def _green_select_modules():
 329     from eventlet.green import select
 330     modules = [('select', select)]
 331
 332     if sys.version_info >= (3, 4):
 333         from eventlet.green import selectors
 334         modules.append(('selectors', selectors))
 335
 336     return modules
 337
 338
 339 def _green_socket_modules():
 340     from eventlet.green import socket
 341     try:
 342         from eventlet.green import ssl
 343         return [('socket', socket), ('ssl', ssl)]
 344     except ImportError:
 345         return [('socket', socket)]
 346
 347
 348 def _green_thread_modules():
 349     from eventlet.green import Queue
 350     from eventlet.green import thread
 351     from eventlet.green import threading
 352     if six.PY2:
 353         return [('Queue', Queue), ('thread', thread), ('threading', threading)]
 354     if six.PY3:
 355         return [('queue', Queue), ('_thread', thread), ('threading', threading)]
 356
 357
 358 def _green_time_modules():
 359     from eventlet.green import time
 360     return [('time', time)]
 361
 362
 363 def _green_MySQLdb():
 364     try:
 365         from eventlet.green import MySQLdb
 366         return [('MySQLdb', MySQLdb)]
 367     except ImportError:
 368         return []
 369
 370
 371 def _green_builtins():
 372     try:
 373         from eventlet.green import builtin
 374         return [('__builtin__' if six.PY2 else 'builtins', builtin)]
 375     except ImportError:
 376         return []
 377
 378
 379 def slurp_properties(source, destination, ignore=[], srckeys=None):
 380     """Copy properties from *source* (assumed to be a module) to
 381     *destination* (assumed to be a dict).
 382
 383     *ignore* lists properties that should not be thusly copied.
 384     *srckeys* is a list of keys to copy, if the source's __all__ is
 385     untrustworthy.
 386     """
 387     if srckeys is None:
 388         srckeys = source.__all__
 389     destination.update(dict([
 390         (name, getattr(source, name))
 391         for name in srckeys
 392         if not (name.startswith('__') or name in ignore)
 393     ]))
 394
 395
 396 if __name__ == "__main__":
 397     sys.argv.pop(0)
 398     monkey_patch()
 399     with open(sys.argv[0]) as f:
 400         code = compile(f.read(), sys.argv[0], 'exec')
 401         exec(code)