python · serhiy-storchaka · Apr 28, 2017 · Apr 28, 2017 · Apr 28, 2017
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -2073,20 +2073,24 @@ 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):
@@ -2100,20 +2104,24 @@ def swap_item(obj, item, new_val):
         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
+            yield real_val
         finally:
             obj[item] = real_val
     else:
         obj[item] = new_val
         try:
             yield
         finally:
-            del obj[item]
+            if item in obj:
+                del obj[item]
 
 def strip_python_stderr(stderr):
     """Strip the stderr of a Python process from potential debug output

diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
@@ -283,17 +283,34 @@ def test_python_is_optimized(self):
 
     def test_swap_attr(self):
         class Obj:
-            x = 1
+            pass
         obj = Obj()
-        with support.swap_attr(obj, "x", 5):
+        obj.x = 1
+        with support.swap_attr(obj, "x", 5) as x:
             self.assertEqual(obj.x, 5)
+            self.assertEqual(x, 1)
         self.assertEqual(obj.x, 1)
+        with support.swap_attr(obj, "y", 5) as y:
+            self.assertEqual(obj.y, 5)
+            self.assertIsNone(y)
+        self.assertFalse(hasattr(obj, 'y'))
+        with support.swap_attr(obj, "y", 5):
+            del obj.y
+        self.assertFalse(hasattr(obj, 'y'))
 
     def test_swap_item(self):
-        D = {"item":1}
-        with support.swap_item(D, "item", 5):
-            self.assertEqual(D["item"], 5)
-        self.assertEqual(D["item"], 1)
+        D = {"x":1}
+        with support.swap_item(D, "x", 5) as x:
+            self.assertEqual(D["x"], 5)
+            self.assertEqual(x, 1)
+        self.assertEqual(D["x"], 1)
+        with support.swap_item(D, "y", 5) as y:
+            self.assertEqual(D["y"], 5)
+            self.assertIsNone(y)
+        self.assertNotIn("y", D)
+        with support.swap_item(D, "y", 5):
+            del D["y"]
+        self.assertNotIn("y", D)
 
     class RefClass:
         attribute1 = None

diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py
@@ -273,13 +273,12 @@ def raise_OSError(*args, **kwargs):
                         tempfile._get_default_tempdir()
                     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(FileNotFoundError):
                         tempfile._get_default_tempdir()

diff --git a/Misc/NEWS b/Misc/NEWS
@@ -199,6 +199,12 @@ Build
 Tests
 -----
 
+- bpo-30197: Enhanced functions swap_attr() and swap_item() in the
+  test.support module.  They now work when delete replaced attribute or item
+  inside the with statement.  The old value of the attribute or item (or None
+  if it doesn't exist) now will be assigned to the target of the "as" clause,
+  if there is one.
+
 - Issue #29571: to match the behaviour of the ``re.LOCALE`` flag,
   test_re.test_locale_flag now uses ``locale.getpreferredencoding(False)`` to
   determine the candidate encoding for the test regex (allowing it to correctly