Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
patcher.py573 linesDownload Raw Back to eventlet
1try:2    import _imp as imp3except ImportError:4    import imp5import sys6try:7    # Only for this purpose, it's irrelevant if `os` was already patched.8    # https://github.com/eventlet/eventlet/pull/6619    from os import register_at_fork10except ImportError:11    register_at_fork = None12 13import eventlet14import six15 16 17__all__ = ['inject', 'import_patched', 'monkey_patch', 'is_monkey_patched']18 19__exclude = set(('__builtins__', '__file__', '__name__'))20 21 22class SysModulesSaver(object):23    """Class that captures some subset of the current state of24    sys.modules.  Pass in an iterator of module names to the25    constructor."""26 27    def __init__(self, module_names=()):28        self._saved = {}29        imp.acquire_lock()30        self.save(*module_names)31 32    def save(self, *module_names):33        """Saves the named modules to the object."""34        for modname in module_names:35            self._saved[modname] = sys.modules.get(modname, None)36 37    def restore(self):38        """Restores the modules that the saver knows about into39        sys.modules.40        """41        try:42            for modname, mod in six.iteritems(self._saved):43                if mod is not None:44                    sys.modules[modname] = mod45                else:46                    try:47                        del sys.modules[modname]48                    except KeyError:49                        pass50        finally:51            imp.release_lock()52 53 54def inject(module_name, new_globals, *additional_modules):55    """Base method for "injecting" greened modules into an imported module.  It56    imports the module specified in *module_name*, arranging things so57    that the already-imported modules in *additional_modules* are used when58    *module_name* makes its imports.59 60    **Note:** This function does not create or change any sys.modules item, so61    if your greened module use code like 'sys.modules["your_module_name"]', you62    need to update sys.modules by yourself.63 64    *new_globals* is either None or a globals dictionary that gets populated65    with the contents of the *module_name* module.  This is useful when creating66    a "green" version of some other module.67 68    *additional_modules* should be a collection of two-element tuples, of the69    form (<name>, <module>).  If it's not specified, a default selection of70    name/module pairs is used, which should cover all use cases but may be71    slower because there are inevitably redundant or unnecessary imports.72    """73    patched_name = '__patched_module_' + module_name74    if patched_name in sys.modules:75        # returning already-patched module so as not to destroy existing76        # references to patched modules77        return sys.modules[patched_name]78 79    if not additional_modules:80        # supply some defaults81        additional_modules = (82            _green_os_modules() +83            _green_select_modules() +84            _green_socket_modules() +85            _green_thread_modules() +86            _green_time_modules())87        # _green_MySQLdb()) # enable this after a short baking-in period88 89    # after this we are gonna screw with sys.modules, so capture the90    # state of all the modules we're going to mess with, and lock91    saver = SysModulesSaver([name for name, m in additional_modules])92    saver.save(module_name)93 94    # Cover the target modules so that when you import the module it95    # sees only the patched versions96    for name, mod in additional_modules:97        sys.modules[name] = mod98 99    # Remove the old module from sys.modules and reimport it while100    # the specified modules are in place101    sys.modules.pop(module_name, None)102    # Also remove sub modules and reimport. Use copy the keys to list103    # because of the pop operations will change the content of sys.modules104    # within th loop105    for imported_module_name in list(sys.modules.keys()):106        if imported_module_name.startswith(module_name + '.'):107            sys.modules.pop(imported_module_name, None)108    try:109        module = __import__(module_name, {}, {}, module_name.split('.')[:-1])110 111        if new_globals is not None:112            # Update the given globals dictionary with everything from this new module113            for name in dir(module):114                if name not in __exclude:115                    new_globals[name] = getattr(module, name)116 117        # Keep a reference to the new module to prevent it from dying118        sys.modules[patched_name] = module119    finally:120        saver.restore()  # Put the original modules back121 122    return module123 124 125def import_patched(module_name, *additional_modules, **kw_additional_modules):126    """Imports a module in a way that ensures that the module uses "green"127    versions of the standard library modules, so that everything works128    nonblockingly.129 130    The only required argument is the name of the module to be imported.131    """132    return inject(133        module_name,134        None,135        *additional_modules + tuple(kw_additional_modules.items()))136 137 138def patch_function(func, *additional_modules):139    """Decorator that returns a version of the function that patches140    some modules for the duration of the function call.  This is141    deeply gross and should only be used for functions that import142    network libraries within their function bodies that there is no143    way of getting around."""144    if not additional_modules:145        # supply some defaults146        additional_modules = (147            _green_os_modules() +148            _green_select_modules() +149            _green_socket_modules() +150            _green_thread_modules() +151            _green_time_modules())152 153    def patched(*args, **kw):154        saver = SysModulesSaver()155        for name, mod in additional_modules:156            saver.save(name)157            sys.modules[name] = mod158        try:159            return func(*args, **kw)160        finally:161            saver.restore()162    return patched163 164 165def _original_patch_function(func, *module_names):166    """Kind of the contrapositive of patch_function: decorates a167    function such that when it's called, sys.modules is populated only168    with the unpatched versions of the specified modules.  Unlike169    patch_function, only the names of the modules need be supplied,170    and there are no defaults.  This is a gross hack; tell your kids not171    to import inside function bodies!"""172    def patched(*args, **kw):173        saver = SysModulesSaver(module_names)174        for name in module_names:175            sys.modules[name] = original(name)176        try:177            return func(*args, **kw)178        finally:179            saver.restore()180    return patched181 182 183def original(modname):184    """ This returns an unpatched version of a module; this is useful for185    Eventlet itself (i.e. tpool)."""186    # note that it's not necessary to temporarily install unpatched187    # versions of all patchable modules during the import of the188    # module; this is because none of them import each other, except189    # for threading which imports thread190    original_name = '__original_module_' + modname191    if original_name in sys.modules:192        return sys.modules.get(original_name)193 194    # re-import the "pure" module and store it in the global _originals195    # dict; be sure to restore whatever module had that name already196    saver = SysModulesSaver((modname,))197    sys.modules.pop(modname, None)198    # some rudimentary dependency checking -- fortunately the modules199    # we're working on don't have many dependencies so we can just do200    # some special-casing here201    if six.PY2:202        deps = {'threading': 'thread', 'Queue': 'threading'}203    if six.PY3:204        deps = {'threading': '_thread', 'queue': 'threading'}205    if modname in deps:206        dependency = deps[modname]207        saver.save(dependency)208        sys.modules[dependency] = original(dependency)209    try:210        real_mod = __import__(modname, {}, {}, modname.split('.')[:-1])211        if modname in ('Queue', 'queue') and not hasattr(real_mod, '_threading'):212            # tricky hack: Queue's constructor in <2.7 imports213            # threading on every instantiation; therefore we wrap214            # it so that it always gets the original threading215            real_mod.Queue.__init__ = _original_patch_function(216                real_mod.Queue.__init__,217                'threading')218        # save a reference to the unpatched module so it doesn't get lost219        sys.modules[original_name] = real_mod220    finally:221        saver.restore()222 223    return sys.modules[original_name]224 225 226already_patched = {}227 228 229def monkey_patch(**on):230    """Globally patches certain system modules to be greenthread-friendly.231 232    The keyword arguments afford some control over which modules are patched.233    If no keyword arguments are supplied, all possible modules are patched.234    If keywords are set to True, only the specified modules are patched.  E.g.,235    ``monkey_patch(socket=True, select=True)`` patches only the select and236    socket modules.  Most arguments patch the single module of the same name237    (os, time, select).  The exceptions are socket, which also patches the ssl238    module if present; and thread, which patches thread, threading, and Queue.239 240    It's safe to call monkey_patch multiple times.241    """242 243    # Workaround for import cycle observed as following in monotonic244    # RuntimeError: no suitable implementation for this system245    # see https://github.com/eventlet/eventlet/issues/401#issuecomment-325015989246    #247    # Make sure the hub is completely imported before any248    # monkey-patching, or we risk recursion if the process of importing249    # the hub calls into monkey-patched modules.250    eventlet.hubs.get_hub()251 252    accepted_args = set(('os', 'select', 'socket',253                         'thread', 'time', 'psycopg', 'MySQLdb',254                         'builtins', 'subprocess'))255    # To make sure only one of them is passed here256    assert not ('__builtin__' in on and 'builtins' in on)257    try:258        b = on.pop('__builtin__')259    except KeyError:260        pass261    else:262        on['builtins'] = b263 264    default_on = on.pop("all", None)265 266    for k in six.iterkeys(on):267        if k not in accepted_args:268            raise TypeError("monkey_patch() got an unexpected "269                            "keyword argument %r" % k)270    if default_on is None:271        default_on = not (True in on.values())272    for modname in accepted_args:273        if modname == 'MySQLdb':274            # MySQLdb is only on when explicitly patched for the moment275            on.setdefault(modname, False)276        if modname == 'builtins':277            on.setdefault(modname, False)278        on.setdefault(modname, default_on)279 280    if on['thread'] and not already_patched.get('thread'):281        _green_existing_locks()282 283    modules_to_patch = []284    for name, modules_function in [285        ('os', _green_os_modules),286        ('select', _green_select_modules),287        ('socket', _green_socket_modules),288        ('thread', _green_thread_modules),289        ('time', _green_time_modules),290        ('MySQLdb', _green_MySQLdb),291        ('builtins', _green_builtins),292        ('subprocess', _green_subprocess_modules),293    ]:294        if on[name] and not already_patched.get(name):295            modules_to_patch += modules_function()296            already_patched[name] = True297 298    if on['psycopg'] and not already_patched.get('psycopg'):299        try:300            from eventlet.support import psycopg2_patcher301            psycopg2_patcher.make_psycopg_green()302            already_patched['psycopg'] = True303        except ImportError:304            # note that if we get an importerror from trying to305            # monkeypatch psycopg, we will continually retry it306            # whenever monkey_patch is called; this should not be a307            # performance problem but it allows is_monkey_patched to308            # tell us whether or not we succeeded309            pass310 311    _threading = original('threading')312    imp.acquire_lock()313    try:314        for name, mod in modules_to_patch:315            orig_mod = sys.modules.get(name)316            if orig_mod is None:317                orig_mod = __import__(name)318            for attr_name in mod.__patched__:319                patched_attr = getattr(mod, attr_name, None)320                if patched_attr is not None:321                    setattr(orig_mod, attr_name, patched_attr)322            deleted = getattr(mod, '__deleted__', [])323            for attr_name in deleted:324                if hasattr(orig_mod, attr_name):325                    delattr(orig_mod, attr_name)326 327            # https://github.com/eventlet/eventlet/issues/592328            if name == 'threading' and register_at_fork:329                def fix_threading_active(330                    _global_dict=_threading.current_thread.__globals__,331                    # alias orig_mod as patched to reflect its new state332                    # https://github.com/eventlet/eventlet/pull/661#discussion_r509877481333                    _patched=orig_mod,334                ):335                    _prefork_active = [None]336 337                    def before_fork():338                        _prefork_active[0] = _global_dict['_active']339                        _global_dict['_active'] = _patched._active340 341                    def after_fork():342                        _global_dict['_active'] = _prefork_active[0]343 344                    register_at_fork(345                        before=before_fork,346                        after_in_parent=after_fork)347                fix_threading_active()348    finally:349        imp.release_lock()350 351    if sys.version_info >= (3, 3):352        import importlib._bootstrap353        thread = original('_thread')354        # importlib must use real thread locks, not eventlet.Semaphore355        importlib._bootstrap._thread = thread356 357        # Issue #185: Since Python 3.3, threading.RLock is implemented in C and358        # so call a C function to get the thread identifier, instead of calling359        # threading.get_ident(). Force the Python implementation of RLock which360        # calls threading.get_ident() and so is compatible with eventlet.361        import threading362        threading.RLock = threading._PyRLock363 364    # Issue #508: Since Python 3.7 queue.SimpleQueue is implemented in C,365    # causing a deadlock.  Replace the C implementation with the Python one.366    if sys.version_info >= (3, 7):367        import queue368        queue.SimpleQueue = queue._PySimpleQueue369 370 371def is_monkey_patched(module):372    """Returns True if the given module is monkeypatched currently, False if373    not.  *module* can be either the module itself or its name.374 375    Based entirely off the name of the module, so if you import a376    module some other way than with the import keyword (including377    import_patched), this might not be correct about that particular378    module."""379    return module in already_patched or \380        getattr(module, '__name__', None) in already_patched381 382 383def _green_existing_locks():384    """Make locks created before monkey-patching safe.385 386    RLocks rely on a Lock and on Python 2, if an unpatched Lock blocks, it387    blocks the native thread. We need to replace these with green Locks.388 389    This was originally noticed in the stdlib logging module."""390    import gc391    import threading392    import eventlet.green.thread393    lock_type = type(threading.Lock())394    rlock_type = type(threading.RLock())395    if hasattr(threading, '_PyRLock'):396        # this happens on CPython3 and PyPy >= 7.0.0: "py3-style" rlocks, they397        # are implemented natively in C and RPython respectively398        py3_style = True399        pyrlock_type = type(threading._PyRLock())400    else:401        # this happens on CPython2.7 and PyPy < 7.0.0: "py2-style" rlocks,402        # they are implemented in pure-python403        py3_style = False404        pyrlock_type = None405 406    # We're monkey-patching so there can't be any greenlets yet, ergo our thread407    # ID is the only valid owner possible.408    tid = eventlet.green.thread.get_ident()409    for obj in gc.get_objects():410        if isinstance(obj, rlock_type):411            if not py3_style and isinstance(obj._RLock__block, lock_type):412                _fix_py2_rlock(obj, tid)413            elif py3_style and not isinstance(obj, pyrlock_type):414                _fix_py3_rlock(obj, tid)415 416    if sys.version_info < (3, 10):417        # Older py3 won't have RLocks show up in gc.get_objects() -- see418        # https://github.com/eventlet/eventlet/issues/546 -- so green a handful419        # that we know are significant420        import logging421        if isinstance(logging._lock, rlock_type):422            _fix_py3_rlock(logging._lock, tid)423        logging._acquireLock()424        try:425            for ref in logging._handlerList:426                handler = ref()427                if handler and isinstance(handler.lock, rlock_type):428                    _fix_py3_rlock(handler.lock, tid)429                del handler430        finally:431            logging._releaseLock()432 433 434def _fix_py2_rlock(rlock, tid):435    import eventlet.green.threading436    old = rlock._RLock__block437    new = eventlet.green.threading.Lock()438    rlock._RLock__block = new439    if old.locked():440        new.acquire()441        rlock._RLock__owner = tid442 443 444def _fix_py3_rlock(old, tid):445    import gc446    import threading447    from eventlet.green.thread import allocate_lock448    new = threading._PyRLock()449    if not hasattr(new, "_block") or not hasattr(new, "_owner"):450        # These will only fail if Python changes its internal implementation of451        # _PyRLock:452        raise RuntimeError(453            "INTERNAL BUG. Perhaps you are using a major version " +454            "of Python that is unsupported by eventlet? Please file a bug " +455            "at https://github.com/eventlet/eventlet/issues/new")456    new._block = allocate_lock()457    acquired = False458    while old._is_owned():459        old.release()460        new.acquire()461        acquired = True462    if old._is_owned():463        new.acquire()464        acquired = True465    if acquired:466        new._owner = tid467    gc.collect()468    for ref in gc.get_referrers(old):469        if isinstance(ref, dict):470            for k, v in list(ref.items()):471                if v is old:472                    ref[k] = new473            continue474        if isinstance(ref, list):475            for i, v in enumerate(ref):476                if v is old:477                    ref[i] = new478            continue479        try:480            ref_vars = vars(ref)481        except TypeError:482            pass483        else:484            for k, v in ref_vars.items():485                if v is old:486                    setattr(ref, k, new)487 488 489def _green_os_modules():490    from eventlet.green import os491    return [('os', os)]492 493 494def _green_select_modules():495    from eventlet.green import select496    modules = [('select', select)]497 498    if sys.version_info >= (3, 4):499        from eventlet.green import selectors500        modules.append(('selectors', selectors))501 502    return modules503 504 505def _green_socket_modules():506    from eventlet.green import socket507    try:508        from eventlet.green import ssl509        return [('socket', socket), ('ssl', ssl)]510    except ImportError:511        return [('socket', socket)]512 513 514def _green_subprocess_modules():515    from eventlet.green import subprocess516    return [('subprocess', subprocess)]517 518 519def _green_thread_modules():520    from eventlet.green import Queue521    from eventlet.green import thread522    from eventlet.green import threading523    if six.PY2:524        return [('Queue', Queue), ('thread', thread), ('threading', threading)]525    if six.PY3:526        return [('queue', Queue), ('_thread', thread), ('threading', threading)]527 528 529def _green_time_modules():530    from eventlet.green import time531    return [('time', time)]532 533 534def _green_MySQLdb():535    try:536        from eventlet.green import MySQLdb537        return [('MySQLdb', MySQLdb)]538    except ImportError:539        return []540 541 542def _green_builtins():543    try:544        from eventlet.green import builtin545        return [('__builtin__' if six.PY2 else 'builtins', builtin)]546    except ImportError:547        return []548 549 550def slurp_properties(source, destination, ignore=[], srckeys=None):551    """Copy properties from *source* (assumed to be a module) to552    *destination* (assumed to be a dict).553 554    *ignore* lists properties that should not be thusly copied.555    *srckeys* is a list of keys to copy, if the source's __all__ is556    untrustworthy.557    """558    if srckeys is None:559        srckeys = source.__all__560    destination.update(dict([561        (name, getattr(source, name))562        for name in srckeys563        if not (name.startswith('__') or name in ignore)564    ]))565 566 567if __name__ == "__main__":568    sys.argv.pop(0)569    monkey_patch()570    with open(sys.argv[0]) as f:571        code = compile(f.read(), sys.argv[0], 'exec')572        exec(code)573 
codekingpro/portable-devtools · Team Ai