1
# Copyright (c) 2010 testtools developers. See LICENSE for details.
3
"""Helpers for monkey-patching Python code."""
11
class MonkeyPatcher(object):
12
"""A set of monkey-patches that can be applied and removed all together.
14
Use this to cover up attributes with new objects. Particularly useful for
15
testing difficult code.
18
# Marker used to indicate that the patched attribute did not exist on the
19
# object before we patched it.
20
_NO_SUCH_ATTRIBUTE = object()
22
def __init__(self, *patches):
23
"""Construct a `MonkeyPatcher`.
25
:param patches: The patches to apply, each should be (obj, name,
26
new_value). Providing patches here is equivalent to calling
29
# List of patches to apply in (obj, name, value).
30
self._patches_to_apply = []
31
# List of the original values for things that have been patched.
32
# (obj, name, value) format.
35
self.add_patch(*patch)
37
def add_patch(self, obj, name, value):
38
"""Add a patch to overwrite 'name' on 'obj' with 'value'.
40
The attribute C{name} on C{obj} will be assigned to C{value} when
41
C{patch} is called or during C{run_with_patches}.
43
You can restore the original values with a call to restore().
45
self._patches_to_apply.append((obj, name, value))
48
"""Apply all of the patches that have been specified with `add_patch`.
50
Reverse this operation using L{restore}.
52
for obj, name, value in self._patches_to_apply:
53
original_value = getattr(obj, name, self._NO_SUCH_ATTRIBUTE)
54
self._originals.append((obj, name, original_value))
55
setattr(obj, name, value)
58
"""Restore all original values to any patched objects.
60
If the patched attribute did not exist on an object before it was
61
patched, `restore` will delete the attribute so as to return the
62
object to its original state.
64
while self._originals:
65
obj, name, value = self._originals.pop()
66
if value is self._NO_SUCH_ATTRIBUTE:
69
setattr(obj, name, value)
71
def run_with_patches(self, f, *args, **kw):
72
"""Run 'f' with the given args and kwargs with all patches applied.
74
Restores all objects to their original state when finished.
83
def patch(obj, attribute, value):
84
"""Set 'obj.attribute' to 'value' and return a callable to restore 'obj'.
86
If 'attribute' is not set on 'obj' already, then the returned callable
87
will delete the attribute when called.
89
:param obj: An object to monkey-patch.
90
:param attribute: The name of the attribute to patch.
91
:param value: The value to set 'obj.attribute' to.
92
:return: A nullary callable that, when run, will restore 'obj' to its
95
patcher = MonkeyPatcher((obj, attribute, value))
97
return patcher.restore