Better support for finalization with weakrefs

Lib/test/test_weakref.py

 import gc
 import sys
 import unittest
 import collections
 import weakref
 import operator
 import contextlib
 import copy
+import subprocess
 
 from test import support
 
 # Used in ReferencesTestCase.test_ref_created_during_del() .
 ref_from_del = None
 
+# Used by FinalizeTestCase as a global that may be replaced by None
+# when the interpreter shuts down.
+_global_var = 'foobar'
+
 class C:
     def method(self):
         pass
 
 
 class Callable:
     bar = None
 
     def __call__(self, x):
         self.bar = x
@@ skipping 1271 matching lines @@
     type2test = weakref.WeakValueDictionary
     def _reference(self):
         return self.__ref.copy()
 
 class WeakKeyDictionaryTestCase(mapping_tests.BasicTestMappingProtocol):
     """Check that WeakKeyDictionary conforms to the mapping protocol"""
     __ref = {Object("key1"):1, Object("key2"):2, Object("key3"):3}
     type2test = weakref.WeakKeyDictionary
     def _reference(self):
         return self.__ref.copy()
+
+
+class FinalizeTestCase(unittest.TestCase):
+
+    class A:
+        pass
+
+    def _collect_if_necessary(self):
+        # we create no ref-cycles so in CPython no gc should be needed
+        if sys.implementation.name != 'cpython':
+            support.gc_collect()
+
+    def test_finalize(self):
+        def add(x,y,z):
+            res.append(x + y + z)
+            return x + y + z
+
+        a = self.A()
+
+        res = []
+        f = weakref.finalize(a, add, 67, 43, z=89)
+        wr, func, args, kwds = f.get()
+        self.assertEqual([wr(), func, args, kwds], [a, add, (67,43), {'z':89}])
+        self.assertEqual(f(), 199)              # deregisters callback
+        self.assertEqual(f(), None)
+        self.assertEqual(f(), None)
+        self.assertEqual(f.get(), None)
+        self.assertEqual(f.pop(), None)
+        self.assertEqual(res, [199])
+
+        res = []
+        f = weakref.finalize(a, add, 67, 43, z=89)
+        wr, func, args, kwds = f.pop()          # deregisters callback
+        self.assertEqual([wr(), func, args, kwds], [a, add, (67,43), {'z':89}])
+        self.assertEqual(f(), None)
+        self.assertEqual(f(), None)
+        self.assertEqual(f.get(), None)
+        self.assertEqual(f.pop(), None)
+        self.assertEqual(res, [])
+
+        res = []
+        f = weakref.finalize(a, add, 67, 43, z=89)
+        del a                                   # deregisters callback
+        self._collect_if_necessary()
+        self.assertEqual(f(), None)
+        self.assertEqual(f(), None)
+        self.assertEqual(f.get(), None)
+        self.assertEqual(f.pop(), None)
+        self.assertEqual(res, [199])
+
+    def test_order(self):
+        a = self.A()
+        res = []
+
+        f1 = weakref.finalize(a, res.append, 'f1')
+        f2 = weakref.finalize(a, res.append, 'f2')
+        f3 = weakref.finalize(a, res.append, 'f3')
+        f4 = weakref.finalize(a, res.append, 'f4')
+        f5 = weakref.finalize(a, res.append, 'f5')
+
+        # make sure finalizers can keep themselves alive
+        del f1, f4
+
+        self.assertTrue(f2.get())
+        self.assertTrue(f3.get())
+        self.assertTrue(f5.get())
+
+        self.assertTrue(f5.pop())
+        self.assertFalse(f5.get())
+
+        f5()                       # nothing because previously unregistered
+        res.append('A')
+        f3()                       # => res.append('f3')
+        self.assertFalse(f3.get())
+        res.append('B')
+        f3()                       # nothing because previously called
+        res.append('C')
+        del a
+        self._collect_if_necessary()
+                                   # => res.append('f4')
+                                   # => res.append('f2')
+                                   # => res.append('f1')
+        self.assertFalse(f2.get())
+        res.append('D')
+        f2()                       # nothing because previously called by gc
+
+        expected = ['A', 'f3', 'B', 'C', 'f4', 'f2', 'f1', 'D']
+        self.assertEqual(res, expected)
+
+    def test_all_freed(self):
+        # we want a weakrefable subclass of weakref.finalize
+        class myfinalize(weakref.finalize):
+            pass
+
+        a = self.A()
+        res = []
+        def callback():
+            res.append(123)
+        f = myfinalize(a, callback)
+
+        wr_callback = weakref.ref(callback)
+        wr_f = weakref.ref(f)
+        del callback, f
+
+        self.assertIsNotNone(wr_callback())
+        self.assertIsNotNone(wr_f())
+
+        del a
+        self._collect_if_necessary()
+
+        self.assertIsNone(wr_callback())
+        self.assertIsNone(wr_f())
+        self.assertEqual(res, [123])
+
+    def test_traceback(self):
+        def f():
+            g()
+
+        def g():
+            1/0
+
+        with support.captured_stderr() as stderr:
+            a = self.A()
+            weakref.finalize(a, f)
+            del a
+            self._collect_if_necessary()
+
+        def process(tb):
+            lines = [line.strip() for line in tb.split('\n')]
+            lines[1] = '...' + lines[1][-6:]
+            lines[3] = '...' + lines[3][-6:]
+            return lines
+
+        expected = '''\
+            Ignored exception in finalizer - Traceback (most recent call last):
+              File ..., line ..., in f
+                g()
+              File ..., line ..., in g
+                1/0
+            ZeroDivisionError: division by zero
+        '''
+        self.assertEqual(process(stderr.getvalue()), process(expected))
+
+    @classmethod
+    def run_in_child(cls):
+        # cls should stay alive till atexit callbacks run
+        f1 = weakref.finalize(cls, print, 'f1', _global_var)
+        f2 = weakref.finalize(cls, print, 'f2', _global_var)
+        f3 = weakref.finalize(cls, print, 'f3', _global_var)
+        f1.atexit = True
+        assert f2.atexit == False
+        f3.atexit = True
+
+    def test_atexit(self):
+        prog = ('from test.test_weakref import FinalizeTestCase;'+
+                'FinalizeTestCase.run_in_child()')
+        output = subprocess.check_output([sys.executable, '-c', prog])

[AntoinePitrou, 2012/08/02 23:19:01]

You should use the helpers from test.script_helper instead.

+        output = output.decode('ascii')
+        output = output.strip().split('\n')
+        self.assertEqual(output, ['f3 foobar', 'f1 foobar'])
+
 
 libreftest = """ Doctest for examples in the library reference: weakref.rst
 
 >>> import weakref
 >>> class Dict(dict):
 ...     pass
 ...
 >>> obj = Dict(red=1, green=2, blue=3)   # this object is weak referencable
 >>> r = weakref.ref(obj)
 >>> print(r() is obj)
@@ skipping 61 matching lines @@
 >>> id2obj(a_id) is a
 True
 >>> del a
 >>> try:
 ...     id2obj(a_id)
 ... except KeyError:
 ...     print('OK')
 ... else:
 ...     print('WeakValueDictionary error')
 OK
+
+
+>>> import weakref

[AntoinePitrou, 2012/08/02 23:19:01]

My personal opinion about doctests is that we shouldn't add any more of them,
since they're tedious to maintain. Also, it doesn't seem they're testing
anything that's not already tested by the unit tests above.

[sbt, 2012/08/03 01:02:14]

I did wonder whether it would be better to make the unittests pass
Doc/library/weakref.rst to doctest.  (It passes.)

+>>> class Object: pass
+...
+>>> kenny = Object()
+>>> weakref.finalize(kenny, print, "you killed kenny!")  #doctest:+ELLIPSIS
+<weakref.finalize object at ...>
+>>> del kenny
+you killed kenny!
+
+
+>>> def callback(x, y, z):
+...    print("CALLBACK")
+...    return x + y + z
+...
+>>> obj = Object()
+>>> f = weakref.finalize(obj, callback, 1, 2, z=3)
+>>> f.get()                                              #doctest:+ELLIPSIS
+(<weakref ... to 'Object' ...>, <function callback ...>, (1, 2), {'z': 3})
+>>> wr, func, args, kwds = _
+>>> assert wr() is obj
+>>> assert func is callback
+
+
+>>> f()
+CALLBACK
+6
+>>> f(), f.get(), f.pop()
+(None, None, None)
+>>> del obj
+
+
+>>> obj = Object()
+>>> f = weakref.finalize(obj, callback, 1, 2, z=3)
+>>> wr, func, args, kwds = f.pop()
+>>> f(), f.get(), f.pop()
+(None, None, None)
+>>> del obj
+>>> func(*args, **kwds)     # explicitly call the popped callback
+CALLBACK
+6
+
+
+>>> kenny = Object()
+>>> f = weakref.finalize(kenny, print, "kenny must die!")
+>>> f.atexit = True
+>>> exit()                                               #doctest:+SKIP
+kenny must die!
 
 """
 
 __test__ = {'libreftest' : libreftest}
 
 def test_main():
     support.run_unittest(
         ReferencesTestCase,
         MappingTestCase,
         WeakValueDictionaryTestCase,
         WeakKeyDictionaryTestCase,
         SubclassableWeakrefTestCase,
+        FinalizeTestCase,
         )
     support.run_doctest(sys.modules[__name__])
 
 
 if __name__ == "__main__":
     test_main()