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