diff options
Diffstat (limited to 'Lib/test/support.py')
| -rw-r--r-- | Lib/test/support.py | 57 |
1 files changed, 56 insertions, 1 deletions
diff --git a/Lib/test/support.py b/Lib/test/support.py index a05f420..28f7e27 100644 --- a/Lib/test/support.py +++ b/Lib/test/support.py @@ -30,7 +30,8 @@ __all__ = ["Error", "TestFailed", "ResourceDenied", "import_module", "run_with_locale", "set_memlimit", "bigmemtest", "bigaddrspacetest", "BasicTestRunner", "run_unittest", "run_doctest", "threading_setup", "threading_cleanup", - "reap_children", "cpython_only", "check_impl_detail", "get_attribute"] + "reap_children", "cpython_only", "check_impl_detail", "get_attribute", + "swap_item", "swap_attr"] class Error(Exception): """Base class for regression test exceptions.""" @@ -1074,3 +1075,57 @@ def reap_children(): break except: break + +@contextlib.contextmanager +def swap_attr(obj, attr, new_val): + """Temporary swap out an attribute with a new object. + + Usage: + with swap_attr(obj, "attr", 5): + ... + + This will set obj.attr to 5 for the duration of the with: block, + restoring the old value at the end of the block. If `attr` doesn't + exist on `obj`, it will be created and then deleted at the end of the + block. + """ + if hasattr(obj, attr): + real_val = getattr(obj, attr) + setattr(obj, attr, new_val) + try: + yield + finally: + setattr(obj, attr, real_val) + else: + setattr(obj, attr, new_val) + try: + yield + finally: + delattr(obj, attr) + +@contextlib.contextmanager +def swap_item(obj, item, new_val): + """Temporary swap out an item with a new object. + + Usage: + with swap_item(obj, "item", 5): + ... + + This will set obj["item"] to 5 for the duration of the with: block, + restoring the old value at the end of the block. If `item` doesn't + exist on `obj`, it will be created and then deleted at the end of the + block. + """ + if item in obj: + real_val = obj[item] + obj[item] = new_val + try: + yield + finally: + obj[item] = real_val + else: + obj[item] = new_val + try: + yield + finally: + del obj[item] |