98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
# Copyright (c) 2010 testtools developers. See LICENSE for details.
|
|
|
|
"""Helpers for monkey-patching Python code."""
|
|
|
|
__all__ = [
|
|
'MonkeyPatcher',
|
|
'patch',
|
|
]
|
|
|
|
|
|
class MonkeyPatcher:
|
|
"""A set of monkey-patches that can be applied and removed all together.
|
|
|
|
Use this to cover up attributes with new objects. Particularly useful for
|
|
testing difficult code.
|
|
"""
|
|
|
|
# Marker used to indicate that the patched attribute did not exist on the
|
|
# object before we patched it.
|
|
_NO_SUCH_ATTRIBUTE = object()
|
|
|
|
def __init__(self, *patches):
|
|
"""Construct a `MonkeyPatcher`.
|
|
|
|
:param patches: The patches to apply, each should be (obj, name,
|
|
new_value). Providing patches here is equivalent to calling
|
|
`add_patch`.
|
|
"""
|
|
# List of patches to apply in (obj, name, value).
|
|
self._patches_to_apply = []
|
|
# List of the original values for things that have been patched.
|
|
# (obj, name, value) format.
|
|
self._originals = []
|
|
for patch in patches:
|
|
self.add_patch(*patch)
|
|
|
|
def add_patch(self, obj, name, value):
|
|
"""Add a patch to overwrite 'name' on 'obj' with 'value'.
|
|
|
|
The attribute C{name} on C{obj} will be assigned to C{value} when
|
|
C{patch} is called or during C{run_with_patches}.
|
|
|
|
You can restore the original values with a call to restore().
|
|
"""
|
|
self._patches_to_apply.append((obj, name, value))
|
|
|
|
def patch(self):
|
|
"""Apply all of the patches that have been specified with `add_patch`.
|
|
|
|
Reverse this operation using L{restore}.
|
|
"""
|
|
for obj, name, value in self._patches_to_apply:
|
|
original_value = getattr(obj, name, self._NO_SUCH_ATTRIBUTE)
|
|
self._originals.append((obj, name, original_value))
|
|
setattr(obj, name, value)
|
|
|
|
def restore(self):
|
|
"""Restore all original values to any patched objects.
|
|
|
|
If the patched attribute did not exist on an object before it was
|
|
patched, `restore` will delete the attribute so as to return the
|
|
object to its original state.
|
|
"""
|
|
while self._originals:
|
|
obj, name, value = self._originals.pop()
|
|
if value is self._NO_SUCH_ATTRIBUTE:
|
|
delattr(obj, name)
|
|
else:
|
|
setattr(obj, name, value)
|
|
|
|
def run_with_patches(self, f, *args, **kw):
|
|
"""Run 'f' with the given args and kwargs with all patches applied.
|
|
|
|
Restores all objects to their original state when finished.
|
|
"""
|
|
self.patch()
|
|
try:
|
|
return f(*args, **kw)
|
|
finally:
|
|
self.restore()
|
|
|
|
|
|
def patch(obj, attribute, value):
|
|
"""Set 'obj.attribute' to 'value' and return a callable to restore 'obj'.
|
|
|
|
If 'attribute' is not set on 'obj' already, then the returned callable
|
|
will delete the attribute when called.
|
|
|
|
:param obj: An object to monkey-patch.
|
|
:param attribute: The name of the attribute to patch.
|
|
:param value: The value to set 'obj.attribute' to.
|
|
:return: A nullary callable that, when run, will restore 'obj' to its
|
|
original state.
|
|
"""
|
|
patcher = MonkeyPatcher((obj, attribute, value))
|
|
patcher.patch()
|
|
return patcher.restore
|