Patch #1680961: remove sys.exitfunc and replace it with a private C API. Also, reimplement atexit in C so it can take advantage of this private API.

Lib/test/test_atexit.py

@@ -4,97 +4,112 @@ import StringIO
 import atexit
 from test import test_support
 
+### helpers
+def h1():
+    print("h1")
+
+def h2():
+    print("h2")
+
+def h3():
+    print("h3")
+
+def h4(*args, **kwargs):
+    print("h4", args, kwargs)
+
+def raise1():
+    raise TypeError
+
+def raise2():
+    raise SystemError
+
 class TestCase(unittest.TestCase):
+    def setUp(self):
+        self.stream = StringIO.StringIO()
+        sys.stdout = sys.stderr = self.stream
+        atexit._clear()
+        
+    def tearDown(self):
+        sys.stdout = sys.__stdout__
+        sys.stderr = sys.__stderr__
+        atexit._clear()
+
     def test_args(self):
         # be sure args are handled properly
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(self.h1)
-            atexit.register(self.h4)
-            atexit.register(self.h4, 4, kw="abc")
-            atexit._run_exitfuncs()
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-        self.assertEqual(s.getvalue(), "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n")
+        atexit.register(h1)
+        atexit.register(h4)
+        atexit.register(h4, 4, kw="abc")
+        atexit._run_exitfuncs()
+
+        self.assertEqual(self.stream.getvalue(),
+                            "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n")
 
     def test_order(self):
         # be sure handlers are executed in reverse order
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(self.h1)
-            atexit.register(self.h2)
-            atexit.register(self.h3)
-            atexit._run_exitfuncs()
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-        self.assertEqual(s.getvalue(), "h3\nh2\nh1\n")
-
-    def test_sys_override(self):
-        # be sure a preset sys.exitfunc is handled properly
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        exfunc = sys.exitfunc
-        sys.exitfunc = self.h1
-        reload(atexit)
-        try:
-            atexit.register(self.h2)
-            atexit._run_exitfuncs()
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-            sys.exitfunc = exfunc
-        self.assertEqual(s.getvalue(), "h2\nh1\n")
+        atexit.register(h1)
+        atexit.register(h2)
+        atexit.register(h3)
+        atexit._run_exitfuncs()
+            
+        self.assertEqual(self.stream.getvalue(), "h3\nh2\nh1\n")
 
     def test_raise(self):
         # be sure raises are handled properly
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(self.raise1)
-            atexit.register(self.raise2)
-            self.assertRaises(TypeError, atexit._run_exitfuncs)
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-
-    ### helpers
-    def h1(self):
-        print("h1")
-
-    def h2(self):
-        print("h2")
-
-    def h3(self):
-        print("h3")
-
-    def h4(self, *args, **kwargs):
-        print("h4", args, kwargs)
-
-    def raise1(self):
-        raise TypeError
-
-    def raise2(self):
-        raise SystemError
+        atexit.register(raise1)
+        atexit.register(raise2)
+        
+        self.assertRaises(TypeError, atexit._run_exitfuncs)
+        
+    def test_stress(self):
+        a = [0]
+        def inc():
+            a[0] += 1
+    
+        for i in range(128):
+            atexit.register(inc)
+        atexit._run_exitfuncs()
+        
+        self.assertEqual(a[0], 128)
+        
+    def test_clear(self):
+        a = [0]
+        def inc():
+            a[0] += 1
+            
+        atexit.register(inc)
+        atexit._clear()
+        atexit._run_exitfuncs()
+        
+        self.assertEqual(a[0], 0)
+        
+    def test_unregister(self):
+        a = [0]
+        def inc():
+            a[0] += 1
+        def dec():
+            a[0] -= 1
+        
+        for i in range(4):    
+            atexit.register(inc)
+        atexit.register(dec)
+        atexit.unregister(inc)
+        atexit._run_exitfuncs()
+        
+        self.assertEqual(a[0], -1)
+        
+    def test_bound_methods(self):
+        l = []
+        atexit.register(l.append, 5)
+        atexit._run_exitfuncs()
+        self.assertEqual(l, [5])
+        
+        atexit.unregister(l.append)
+        atexit._run_exitfuncs()
+        self.assertEqual(l, [5])
+        
 
 def test_main():
     test_support.run_unittest(TestCase)
 
-
 if __name__ == "__main__":
     test_main()