def patch(raise_errors=True, **patch_modules):
"""Patch only a set of given modules.
:param bool raise_errors: Raise error if one patch fail.
:param dict \**patch_modules: List of modules to patch.
>>> patch(psycopg=True, elasticsearch=True)
"""
modules = [
m for (m, should_patch) in patch_modules.items() if should_patch
]
for module in modules:
if module in _PATCH_ON_IMPORT:
# If the module has already been imported then patch immediately
if module in sys.modules:
patch_module(module, raise_errors=raise_errors)
# Otherwise, add a hook to patch when it is imported for the first time
else:
# Use factory to create handler to close over `module` and `raise_errors` values from this loop
when_imported(module)(_on_import_factory(module, raise_errors))
# manually add module to patched modules
_PATCHED_MODULES.add(module)
else:
patch_module(module, raise_errors=raise_errors)
patched_modules = get_patched_modules()
log.info("patched %s/%s modules (%s)", len(patched_modules), len(modules),
",".join(patched_modules))
def _ensure_patch_requests():
"""
`requests` is third-party, may not be installed or used,
but ensure it gets patched if installed and used.
"""
if "requests" in sys.modules:
# already imported, patch now
_patch_requests(sys.modules["requests"])
else:
# patch when imported
when_imported("requests")(_patch_requests)
def instrument(**instrument_libs):
libs = [
l for (l, will_instrument) in instrument_libs.items()
if will_instrument
]
for lib in libs:
if lib in sys.modules:
instrument_lib(lib)
else:
when_imported(lib)(_on_import_wrapper(lib))
_INSTRUMENTED_LIBS.add(lib)
patched_libs = get_already_instrumented()
logger.info("Instrumented {}/{} libraries ({})".format(
len(patched_libs), len(libs), ", ".join(patched_libs)))
def instrument(instrument_libs):
for lib in instrument_libs:
if will_instrument(lib):
if lib in sys.modules:
instrument_lib(lib)
_INSTRUMENTED_LIBS.add(lib)
else:
when_imported(lib)(_on_import_wrapper(lib))
_INSTRUMENT_LIBS_LAZY.add(lib)
patched_libs = get_already_instrumented()
lazy_libs = get_will_instrument()
logger.info(
"Instrumented {}/{} libraries ({}). Will instrument when imported: ({})"
.format(len(patched_libs), len(instrument_libs),
", ".join(patched_libs), ", ".join(lazy_libs)))