merge

2 files changed, 41 insertions, 10 deletions

diff --git a/Lib/test/support.py b/Lib/test/support.py
index f5a53ca628..98f333e289 100644
--- a/Lib/test/support.py
+++ b/Lib/test/support.py
@@ -1478,3 +1478,36 @@ def skip_unless_symlink(test):
     ok = can_symlink()
     msg = "Requires functional symlink implementation"
     return test if ok else unittest.skip(msg)(test)
+
+def patch(test_instance, object_to_patch, attr_name, new_value):
+    """Override 'object_to_patch'.'attr_name' with 'new_value'.
+
+    Also, add a cleanup procedure to 'test_instance' to restore
+    'object_to_patch' value for 'attr_name'.
+    The 'attr_name' should be a valid attribute for 'object_to_patch'.
+
+    """
+    # check that 'attr_name' is a real attribute for 'object_to_patch'
+    # will raise AttributeError if it does not exist
+    getattr(object_to_patch, attr_name)
+
+    # keep a copy of the old value
+    attr_is_local = False
+    try:
+        old_value = object_to_patch.__dict__[attr_name]
+    except (AttributeError, KeyError):
+        old_value = getattr(object_to_patch, attr_name, None)
+    else:
+        attr_is_local = True
+
+    # restore the value when the test is done
+    def cleanup():
+        if attr_is_local:
+            setattr(object_to_patch, attr_name, old_value)
+        else:
+            delattr(object_to_patch, attr_name)
+
+    test_instance.addCleanup(cleanup)
+
+    # actually override the attribute
+    setattr(object_to_patch, attr_name, new_value)
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index c3b2a4fc72..643e2e66eb 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -265,28 +265,26 @@ class DisTests(unittest.TestCase):
         self.do_disassembly_test(method_bytecode, dis_c_instance_method_bytes)
 
     def test_dis_none(self):
+        try:
+            del sys.last_traceback
+        except AttributeError:
+            pass
         self.assertRaises(RuntimeError, dis.dis, None)
 
     def test_dis_object(self):
         self.assertRaises(TypeError, dis.dis, object())
 
     def test_dis_traceback(self):
-        not_defined = object()
-        tb = None
-        old = getattr(sys, 'last_traceback', not_defined)
-
-        def cleanup():
-            if old != not_defined:
-                sys.last_traceback = old
-            else:
-                del sys.last_traceback
+        try:
+            del sys.last_traceback
+        except AttributeError:
+            pass
 
         try:
             1/0
         except Exception as e:
             tb = e.__traceback__
             sys.last_traceback = tb
-            self.addCleanup(cleanup)
 
         tb_dis = self.get_disassemble_as_string(tb.tb_frame.f_code, tb.tb_lasti)
         self.do_disassembly_test(None, tb_dis)