When testing, it can be really handy to temporarily override an object's attribute with a mock object. This recipe provides a context manager that does that.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
import contextlib @contextlib.contextmanager def temp_setattr(ob, attr, new_value): """Temporarily set an attribute on an object for the duration of the context manager.""" replaced = False old_value = None if hasattr(ob, attr): try: if attr in ob.__dict__: replaced = True except AttributeError: if attr in ob.__slots__: replaced = True if replaced: old_value = getattr(ob, attr) setattr(ob, attr, new_value) yield replaced, old_value if not replaced: delattr(ob, attr) else: setattr(ob, attr, old_value)
What separates this recipe from other ones that focus on temporary attribute setting is that it is generic to any object and the care taken to make sure that a shadowed attribute ends up back in its original state. If one blindly used
setattr() if an attribute is found then you can easily end up setting an attribute on an instance that was never there in the first place; the attribute was originally found on an inherited class. This recipe does as much as is reasonably possible to make sure that when the context manager exits everything goes back to the way it was in terms of the state of the object, not just semantically from the perspective of the object.