diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py index e339da5ff4d1d3..976f301c0a9144 100644 --- a/Lib/test/test_support.py +++ b/Lib/test/test_support.py @@ -1647,12 +1647,15 @@ def swap_attr(obj, attr, new_val): 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. + + The old value (or None if it doesn't exist) will be assigned to the + target of the "as" clause, if there is one. """ if hasattr(obj, attr): real_val = getattr(obj, attr) setattr(obj, attr, new_val) try: - yield + yield real_val finally: setattr(obj, attr, real_val) else: @@ -1660,7 +1663,39 @@ def swap_attr(obj, attr, new_val): try: yield finally: - delattr(obj, attr) + if hasattr(obj, attr): + 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. + + The old value (or None if it doesn't exist) will be assigned to the + target of the "as" clause, if there is one. + """ + if item in obj: + real_val = obj[item] + obj[item] = new_val + try: + yield real_val + finally: + obj[item] = real_val + else: + obj[item] = new_val + try: + yield + finally: + if item in obj: + del obj[item] def py3k_bytes(b): """Emulate the py3k bytes() constructor. diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py index 078e4a9a2d06a9..5c85cc9bd306d4 100644 --- a/Lib/test/test_tempfile.py +++ b/Lib/test/test_tempfile.py @@ -235,13 +235,12 @@ def raise_OSError(*args, **kwargs): self.assertEqual(cm.exception.errno, errno.ENOENT) self.assertEqual(os.listdir(our_temp_directory), []) - open = io.open def bad_writer(*args, **kwargs): - fp = open(*args, **kwargs) + fp = orig_open(*args, **kwargs) fp.write = raise_OSError return fp - with support.swap_attr(io, "open", bad_writer): + with support.swap_attr(io, "open", bad_writer) as orig_open: # test again with failing write() with self.assertRaises(IOError) as cm: tempfile._get_default_tempdir() diff --git a/Misc/NEWS b/Misc/NEWS index 657b2fb84e1305..8ec9f2a0eb167b 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -144,6 +144,12 @@ Build Tests ----- +- bpo-30197: Enhanced function swap_attr() in the test.test_support module. + It now works when delete replaced attribute inside the with statement. The + old value of the attribute (or None if it doesn't exist) now will be + assigned to the target of the "as" clause, if there is one. + Also backported function swap_item(). + - bpo-28087: Skip test_asyncore and test_eintr poll failures on macOS. Skip some tests of select.poll when running on macOS due to unresolved issues with the underlying system poll function on some macOS versions.