253 lines
8.8 KiB
Python
253 lines
8.8 KiB
Python
|
# Copyright 2018 gevent. See LICENSE for details.
|
||
|
|
|
||
|
|
# Portions of the following are inspired by code from eventlet. I
|
||
|
|
# believe they are distinct enough that no eventlet copyright would
|
||
|
|
# apply (they are not a copy or substantial portion of the eventlot
|
||
|
|
# code).
|
||
|
|
|
||
|
|
# Added in gevent 1.3a2. Not public in that release.
|
||
|
|
|
||
|
|
from __future__ import absolute_import, print_function
|
||
|
|
|
||
|
|
import importlib
|
||
|
|
import sys
|
||
|
|
|
||
|
|
|
||
|
|
from gevent._compat import iteritems
|
||
|
|
from gevent._compat import imp_acquire_lock
|
||
|
|
from gevent._compat import imp_release_lock
|
||
|
|
|
||
|
|
|
||
|
|
from gevent.builtins import __import__ as g_import
|
||
|
|
|
||
|
|
|
||
|
|
MAPPING = {
|
||
|
|
'gevent.local': '_threading_local',
|
||
|
|
'gevent.socket': 'socket',
|
||
|
|
'gevent.select': 'select',
|
||
|
|
'gevent.selectors': 'selectors',
|
||
|
|
'gevent.ssl': 'ssl',
|
||
|
|
'gevent.thread': '_thread',
|
||
|
|
'gevent.subprocess': 'subprocess',
|
||
|
|
'gevent.os': 'os',
|
||
|
|
'gevent.threading': 'threading',
|
||
|
|
'gevent.builtins': 'builtins',
|
||
|
|
'gevent.signal': 'signal',
|
||
|
|
'gevent.time': 'time',
|
||
|
|
'gevent.queue': 'queue',
|
||
|
|
'gevent.contextvars': 'contextvars',
|
||
|
|
}
|
||
|
|
|
||
|
|
OPTIONAL_STDLIB_MODULES = frozenset()
|
||
|
|
_PATCH_PREFIX = '__g_patched_module_'
|
||
|
|
|
||
|
|
def _collect_stdlib_gevent_modules():
|
||
|
|
"""
|
||
|
|
Return a map from standard library name to
|
||
|
|
imported gevent module that provides the same API.
|
||
|
|
|
||
|
|
Optional modules are skipped if they cannot be imported.
|
||
|
|
"""
|
||
|
|
result = {}
|
||
|
|
|
||
|
|
for gevent_name, stdlib_name in iteritems(MAPPING):
|
||
|
|
try:
|
||
|
|
result[stdlib_name] = importlib.import_module(gevent_name)
|
||
|
|
except ImportError:
|
||
|
|
if stdlib_name in OPTIONAL_STDLIB_MODULES:
|
||
|
|
continue
|
||
|
|
raise
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
class _SysModulesPatcher(object):
|
||
|
|
|
||
|
|
def __init__(self, importing, extra_all=lambda mod_name: ()):
|
||
|
|
# Permanent state.
|
||
|
|
self.extra_all = extra_all
|
||
|
|
self.importing = importing
|
||
|
|
# green modules, replacing regularly imported modules.
|
||
|
|
# This begins as the gevent list of modules, and
|
||
|
|
# then gets extended with green things from the tree we import.
|
||
|
|
self._green_modules = _collect_stdlib_gevent_modules()
|
||
|
|
|
||
|
|
## Transient, reset each time we're called.
|
||
|
|
# The set of things imported before we began.
|
||
|
|
self._t_modules_to_restore = {}
|
||
|
|
|
||
|
|
def _save(self):
|
||
|
|
self._t_modules_to_restore = {}
|
||
|
|
|
||
|
|
# Copy all the things we know we are going to overwrite.
|
||
|
|
for modname in self._green_modules:
|
||
|
|
self._t_modules_to_restore[modname] = sys.modules.get(modname, None)
|
||
|
|
|
||
|
|
# Copy anything else in the import tree.
|
||
|
|
for modname, mod in list(iteritems(sys.modules)):
|
||
|
|
if modname.startswith(self.importing):
|
||
|
|
self._t_modules_to_restore[modname] = mod
|
||
|
|
# And remove it. If it had been imported green, it will
|
||
|
|
# be put right back. Otherwise, it was imported "manually"
|
||
|
|
# outside this process and isn't green.
|
||
|
|
del sys.modules[modname]
|
||
|
|
|
||
|
|
# Cover the target modules so that when you import the module it
|
||
|
|
# sees only the patched versions
|
||
|
|
for name, mod in iteritems(self._green_modules):
|
||
|
|
sys.modules[name] = mod
|
||
|
|
|
||
|
|
def _restore(self):
|
||
|
|
# Anything from the same package tree we imported this time
|
||
|
|
# needs to be saved so we can restore it later, and so it doesn't
|
||
|
|
# leak into the namespace.
|
||
|
|
|
||
|
|
for modname, mod in list(iteritems(sys.modules)):
|
||
|
|
if modname.startswith(self.importing):
|
||
|
|
self._green_modules[modname] = mod
|
||
|
|
del sys.modules[modname]
|
||
|
|
|
||
|
|
# Now, what we saved at the beginning needs to be restored.
|
||
|
|
for modname, mod in iteritems(self._t_modules_to_restore):
|
||
|
|
if mod is not None:
|
||
|
|
sys.modules[modname] = mod
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
del sys.modules[modname]
|
||
|
|
except KeyError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def __exit__(self, t, v, tb):
|
||
|
|
try:
|
||
|
|
self._restore()
|
||
|
|
finally:
|
||
|
|
imp_release_lock()
|
||
|
|
self._t_modules_to_restore = None
|
||
|
|
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
imp_acquire_lock()
|
||
|
|
self._save()
|
||
|
|
return self
|
||
|
|
|
||
|
|
module = None
|
||
|
|
|
||
|
|
def __call__(self, after_import_hook):
|
||
|
|
if self.module is None:
|
||
|
|
with self:
|
||
|
|
self.module = self.import_one(self.importing, after_import_hook)
|
||
|
|
# Circular reference. Someone must keep a reference to this module alive
|
||
|
|
# for it to be visible. We record it in sys.modules to be that someone, and
|
||
|
|
# to aid debugging. In the past, we worked with multiple completely separate
|
||
|
|
# invocations of `import_patched`, but we no longer do.
|
||
|
|
self.module.__gevent_patcher__ = self
|
||
|
|
sys.modules[_PATCH_PREFIX + self.importing] = self.module
|
||
|
|
return self
|
||
|
|
|
||
|
|
def import_one(self, module_name, after_import_hook):
|
||
|
|
patched_name = _PATCH_PREFIX + module_name
|
||
|
|
if patched_name in sys.modules:
|
||
|
|
return sys.modules[patched_name]
|
||
|
|
|
||
|
|
assert module_name.startswith(self.importing)
|
||
|
|
sys.modules.pop(module_name, None)
|
||
|
|
|
||
|
|
module = g_import(module_name, {}, {}, module_name.split('.')[:-1])
|
||
|
|
self.module = module
|
||
|
|
# On Python 3, we could probably do something much nicer with the
|
||
|
|
# import machinery? Set the __loader__ or __finder__ or something like that?
|
||
|
|
self._import_all([module])
|
||
|
|
after_import_hook(module)
|
||
|
|
return module
|
||
|
|
|
||
|
|
def _import_all(self, queue):
|
||
|
|
# Called while monitoring for patch changes.
|
||
|
|
while queue:
|
||
|
|
module = queue.pop(0)
|
||
|
|
name = module.__name__
|
||
|
|
mod_all = tuple(getattr(module, '__all__', ())) + self.extra_all(name)
|
||
|
|
for attr_name in mod_all:
|
||
|
|
try:
|
||
|
|
getattr(module, attr_name)
|
||
|
|
except AttributeError:
|
||
|
|
module_name = module.__name__ + '.' + attr_name
|
||
|
|
new_module = g_import(module_name, {}, {}, attr_name)
|
||
|
|
setattr(module, attr_name, new_module)
|
||
|
|
queue.append(new_module)
|
||
|
|
|
||
|
|
|
||
|
|
def import_patched(module_name,
|
||
|
|
extra_all=lambda mod_name: (),
|
||
|
|
after_import_hook=lambda module: None):
|
||
|
|
"""
|
||
|
|
Import *module_name* with gevent monkey-patches active,
|
||
|
|
and return an object holding the greened module as *module*.
|
||
|
|
|
||
|
|
Any sub-modules that were imported by the package are also
|
||
|
|
saved.
|
||
|
|
|
||
|
|
.. versionchanged:: 1.5a4
|
||
|
|
If the module defines ``__all__``, then each of those
|
||
|
|
attributes/modules is also imported as part of the same transaction,
|
||
|
|
recursively. The order of ``__all__`` is respected. Anything passed in
|
||
|
|
*extra_all* (which must be in the same namespace tree) is also imported.
|
||
|
|
|
||
|
|
.. versionchanged:: 1.5a4
|
||
|
|
You must now do all patching for a given module tree
|
||
|
|
with one call to this method, or at least by using the returned
|
||
|
|
object.
|
||
|
|
"""
|
||
|
|
|
||
|
|
with cached_platform_architecture():
|
||
|
|
# Save the current module state, and restore on exit,
|
||
|
|
# capturing desirable changes in the modules package.
|
||
|
|
patcher = _SysModulesPatcher(module_name, extra_all)
|
||
|
|
patcher(after_import_hook)
|
||
|
|
return patcher
|
||
|
|
|
||
|
|
|
||
|
|
class cached_platform_architecture(object):
|
||
|
|
"""
|
||
|
|
Context manager that caches ``platform.architecture``.
|
||
|
|
|
||
|
|
Some things that load shared libraries (like Cryptodome, via
|
||
|
|
dnspython) invoke ``platform.architecture()`` for each one. That
|
||
|
|
in turn wants to fork and run commands , which in turn wants to
|
||
|
|
call ``threading._after_fork`` if the GIL has been initialized.
|
||
|
|
All of that means that certain imports done early may wind up
|
||
|
|
wanting to have the hub initialized potentially much earlier than
|
||
|
|
before.
|
||
|
|
|
||
|
|
Part of the fix is to observe when that happens and delay
|
||
|
|
initializing parts of gevent until as late as possible (e.g., we
|
||
|
|
delay importing and creating the resolver until the hub needs it,
|
||
|
|
unless explicitly configured).
|
||
|
|
|
||
|
|
The rest of the fix is to avoid the ``_after_fork`` issues by
|
||
|
|
first caching the results of platform.architecture before doing
|
||
|
|
patched imports.
|
||
|
|
|
||
|
|
(See events.py for similar issues with platform, and
|
||
|
|
test__threading_2.py for notes about threading._after_fork if the
|
||
|
|
GIL has been initialized)
|
||
|
|
"""
|
||
|
|
|
||
|
|
_arch_result = None
|
||
|
|
_orig_arch = None
|
||
|
|
_platform = None
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
import platform
|
||
|
|
self._platform = platform
|
||
|
|
self._arch_result = platform.architecture()
|
||
|
|
self._orig_arch = platform.architecture
|
||
|
|
def arch(*args, **kwargs):
|
||
|
|
if not args and not kwargs:
|
||
|
|
return self._arch_result
|
||
|
|
return self._orig_arch(*args, **kwargs)
|
||
|
|
platform.architecture = arch
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *_args):
|
||
|
|
self._platform.architecture = self._orig_arch
|
||
|
|
self._platform = None
|