lib/testtools/testtools/monkey.py

   1 # Copyright (c) 2010 testtools developers. See LICENSE for details.
   2
   3 """Helpers for monkey-patching Python code."""
   4
   5 __all__ = [
   6     'MonkeyPatcher',
   7     'patch',
   8     ]
   9
  10
  11 class MonkeyPatcher(object):
  12     """A set of monkey-patches that can be applied and removed all together.
  13
  14     Use this to cover up attributes with new objects. Particularly useful for
  15     testing difficult code.
  16     """
  17
  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()
  21
  22     def __init__(self, *patches):
  23         """Construct a `MonkeyPatcher`.
  24
  25         :param patches: The patches to apply, each should be (obj, name,
  26             new_value). Providing patches here is equivalent to calling
  27             `add_patch`.
  28         """
  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.
  33         self._originals = []
  34         for patch in patches:
  35             self.add_patch(*patch)
  36
  37     def add_patch(self, obj, name, value):
  38         """Add a patch to overwrite 'name' on 'obj' with 'value'.
  39
  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}.
  42
  43         You can restore the original values with a call to restore().
  44         """
  45         self._patches_to_apply.append((obj, name, value))
  46
  47     def patch(self):
  48         """Apply all of the patches that have been specified with `add_patch`.
  49
  50         Reverse this operation using L{restore}.
  51         """
  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)
  56
  57     def restore(self):
  58         """Restore all original values to any patched objects.
  59
  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.
  63         """
  64         while self._originals:
  65             obj, name, value = self._originals.pop()
  66             if value is self._NO_SUCH_ATTRIBUTE:
  67                 delattr(obj, name)
  68             else:
  69                 setattr(obj, name, value)
  70
  71     def run_with_patches(self, f, *args, **kw):
  72         """Run 'f' with the given args and kwargs with all patches applied.
  73
  74         Restores all objects to their original state when finished.
  75         """
  76         self.patch()
  77         try:
  78             return f(*args, **kw)
  79         finally:
  80             self.restore()
  81
  82
  83 def patch(obj, attribute, value):
  84     """Set 'obj.attribute' to 'value' and return a callable to restore 'obj'.
  85
  86     If 'attribute' is not set on 'obj' already, then the returned callable
  87     will delete the attribute when called.
  88
  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
  93         original state.
  94     """
  95     patcher = MonkeyPatcher((obj, attribute, value))
  96     patcher.patch()
  97     return patcher.restore