python · serhiy-storchaka · Apr 28, 2017 · Apr 28, 2017 · Apr 28, 2017
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
@@ -1647,20 +1647,55 @@ 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:
         setattr(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
@@ -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
@@ -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.