diff -r 60f7719c0415 Include/frameobject.h --- a/Include/frameobject.h Fri Apr 13 21:05:36 2012 -0600 +++ b/Include/frameobject.h Sat May 05 16:14:32 2012 +0300 @@ -45,6 +45,7 @@ PyCode_Addr2Line to calculate the line from the current bytecode index. */ int f_lineno; /* Current line number */ + int f_cleanup_num; /* Number of currently nested cleanup scopes */ int f_iblock; /* index in f_blockstack */ PyTryBlock f_blockstack[CO_MAXBLOCKS]; /* for try and loop blocks */ PyObject *f_localsplus[1]; /* locals+stack, dynamically sized */ diff -r 60f7719c0415 Include/pystate.h --- a/Include/pystate.h Fri Apr 13 21:05:36 2012 -0600 +++ b/Include/pystate.h Sat May 05 16:14:32 2012 +0300 @@ -87,6 +87,7 @@ Py_tracefunc c_tracefunc; PyObject *c_profileobj; PyObject *c_traceobj; + PyObject *c_cleanupobj; PyObject *curexc_type; PyObject *curexc_value; diff -r 60f7719c0415 Lib/test/test_cleanuphook.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/Lib/test/test_cleanuphook.py Sat May 05 16:14:32 2012 +0300 @@ -0,0 +1,140 @@ +import sys +import unittest +from contextlib import contextmanager + +from test.support import run_unittest + +#-------------------- Target functions for hooking into ----------------------# +# + +# Generators + +def generator1(): + """Simple generator""" + try: + yield 0 + finally: + yield 1 + yield 0 + +def generator2(): + """Return from try""" + try: + yield 0 + return + finally: + yield 1 + yield 0 + + +def generator3(): + """Return from try without yield""" + yield 0 + try: + return + finally: + yield 1 + yield 0 + + +def generator4(): + """Cascade try""" + yield 0 + try: + yield 0 + finally: + try: + yield 1 + finally: + yield 2 + yield 1 + yield 0 + + +# Context managers + +@contextmanager +def context_manager(test): + test.assertTrue(sys._getframe(2).f_in_cleanup) + try: + yield + finally: + test.assertTrue(sys._getframe(0).f_in_cleanup) + test.assertTrue(sys._getframe(2).f_in_cleanup) + test.assertFalse(sys._getframe(0).f_in_cleanup) + test.assertTrue(sys._getframe(2).f_in_cleanup) + +def context_call(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + with context_manager(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + test.assertFalse(sys._getframe(0).f_in_cleanup) + +def context_call2(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + with context_manager(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + test.assertFalse(sys._getframe(0).f_in_cleanup) + with context_manager(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + test.assertFalse(sys._getframe(0).f_in_cleanup) + +def context_return(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + with context_manager(test): + test.assertFalse(sys._getframe(0).f_in_cleanup) + return + + +#------------------------------ Test cases -----------------------------------# + +class TestCleanupHook(unittest.TestCase): + """White-box testing of line-counting, via runfunc""" + def setUp(self): + self.addCleanup(sys.setcleanuphook, sys.getcleanuphook()) + sys.setcleanuphook(self._hook) + self.hooked = 0 + + def _hook(self, frame): + self.hooked += 1 + + def _test_generator(self, gen): + gen = gen() + for val in gen: + self.assertEquals(val, gen.gi_frame.f_in_cleanup) + + def test_generator1(self): + self._test_generator(generator1) + self.assertEqual(self.hooked, 1) + + def test_generator2(self): + self._test_generator(generator2) + self.assertEqual(self.hooked, 1) + + def test_generator3(self): + self._test_generator(generator3) + self.assertEqual(self.hooked, 1) + + def test_generator4(self): + self._test_generator(generator4) + self.assertEqual(self.hooked, 1) + + def test_context_manager(self): + context_call(self) + self.assertEqual(self.hooked, 3) + + def test_context_manager2(self): + context_call2(self) + self.assertEqual(self.hooked, 6) + + def test_context_manager_return(self): + context_return(self) + self.assertEqual(self.hooked, 3) + + +def test_main(): + run_unittest(__name__) + + +if __name__ == '__main__': + test_main() diff -r 60f7719c0415 Objects/frameobject.c --- a/Objects/frameobject.c Fri Apr 13 21:05:36 2012 -0600 +++ b/Objects/frameobject.c Sat May 05 16:14:32 2012 +0300 @@ -15,11 +15,12 @@ #define OFF(x) offsetof(PyFrameObject, x) static PyMemberDef frame_memberlist[] = { - {"f_back", T_OBJECT, OFF(f_back), READONLY}, - {"f_code", T_OBJECT, OFF(f_code), READONLY}, - {"f_builtins", T_OBJECT, OFF(f_builtins), READONLY}, - {"f_globals", T_OBJECT, OFF(f_globals), READONLY}, - {"f_lasti", T_INT, OFF(f_lasti), READONLY}, + {"f_back", T_OBJECT, OFF(f_back), READONLY}, + {"f_code", T_OBJECT, OFF(f_code), READONLY}, + {"f_builtins", T_OBJECT, OFF(f_builtins), READONLY}, + {"f_globals", T_OBJECT, OFF(f_globals), READONLY}, + {"f_lasti", T_INT, OFF(f_lasti), READONLY}, + {"f_in_cleanup", T_INT, OFF(f_cleanup_num), READONLY}, {NULL} /* Sentinel */ }; @@ -711,6 +712,7 @@ f->f_lasti = -1; f->f_lineno = code->co_firstlineno; f->f_iblock = 0; + f->f_cleanup_num = 0; _PyObject_GC_TRACK(f); return f; diff -r 60f7719c0415 Python/ceval.c --- a/Python/ceval.c Fri Apr 13 21:05:36 2012 -0600 +++ b/Python/ceval.c Sat May 05 16:14:32 2012 +0300 @@ -1891,12 +1891,16 @@ { PyTryBlock *b = PyFrame_BlockPop(f); UNWIND_BLOCK(b); + if(b->b_type == SETUP_FINALLY) { + f->f_cleanup_num ++; + } } DISPATCH(); PREDICTED(END_FINALLY); TARGET(END_FINALLY) v = POP(); + f->f_cleanup_num --; if (PyLong_Check(v)) { why = (enum why_code) PyLong_AS_LONG(v); assert(why != WHY_YIELD); @@ -1919,6 +1923,10 @@ u = POP(); PyErr_Restore(v, w, u); why = WHY_RERAISE; + if(!f->f_cleanup_num && tstate->c_cleanupobj) { + PyObject_CallFunctionObjArgs(tstate->c_cleanupobj, f, + NULL); + } break; } else if (v != Py_None) { @@ -1927,6 +1935,11 @@ why = WHY_EXCEPTION; } Py_DECREF(v); + + if(!f->f_cleanup_num && tstate->c_cleanupobj) { + PyObject_CallFunctionObjArgs(tstate->c_cleanupobj, f, NULL); + } + break; TARGET(LOAD_BUILD_CLASS) @@ -2551,8 +2564,18 @@ x = NULL; break; } + f->f_cleanup_num ++; x = PyObject_CallFunctionObjArgs(u, NULL); + f->f_cleanup_num --; Py_DECREF(u); + + if(!f->f_cleanup_num && tstate->c_cleanupobj) { + u = PyObject_CallFunctionObjArgs(tstate->c_cleanupobj, f, + NULL); + if (!u) + break; + } + if (!x) break; /* Setup the finally block before pushing the result @@ -2641,8 +2664,12 @@ x = PyObject_CallFunctionObjArgs(exit_func, u, v, w, NULL); Py_DECREF(exit_func); - if (x == NULL) + + + if (x == NULL) { + f->f_cleanup_num --; break; /* Go to error exit */ + } if (u != Py_None) err = PyObject_IsTrue(x); @@ -2650,9 +2677,10 @@ err = 0; Py_DECREF(x); - if (err < 0) + if (err < 0) { + f->f_cleanup_num --; break; /* Go to error exit */ - else if (err > 0) { + } else if (err > 0) { err = 0; /* There was an exception and a True return */ PUSH(PyLong_FromLong((long) WHY_SILENCED)); @@ -2984,6 +3012,7 @@ PUSH(val); PUSH(exc); why = WHY_NOT; + f->f_cleanup_num ++; JUMPTO(handler); break; } @@ -2992,6 +3021,7 @@ PUSH(retval); PUSH(PyLong_FromLong((long)why)); why = WHY_NOT; + f->f_cleanup_num ++; JUMPTO(b->b_handler); break; } diff -r 60f7719c0415 Python/sysmodule.c --- a/Python/sysmodule.c Fri Apr 13 21:05:36 2012 -0600 +++ b/Python/sysmodule.c Sat May 05 16:14:32 2012 +0300 @@ -451,6 +451,55 @@ ); static PyObject * +sys_setcleanuphook(PyObject *self, PyObject *args) +{ + PyObject *hook; + if(!PyArg_ParseTuple(args, "O", &hook)) { + return NULL; + } + Py_INCREF(hook); + + PyThreadState *tstate = PyThreadState_GET(); + PyObject *oldhook = tstate->c_cleanupobj; + Py_XDECREF(oldhook); + + if(hook == Py_None) { + tstate->c_cleanupobj = NULL; + Py_DECREF(hook); + } else { + tstate->c_cleanupobj = hook; + } + + Py_INCREF(Py_None); + return Py_None; +} + +PyDoc_STRVAR(setcleanuphook_doc, +"setcleanuphook(function)\n\ +\n\ +Set the global cleanup hook. The function is called when ``f_in_cleanup``\n\ +counter drops to zero." +); + +static PyObject * +sys_getcleanuphook(PyObject *self, PyObject *args) +{ + PyThreadState *tstate = PyThreadState_GET(); + PyObject *temp = tstate->c_cleanupobj; + + if (temp == NULL) + temp = Py_None; + Py_INCREF(temp); + return temp; +} + +PyDoc_STRVAR(getcleanuphook_doc, +"getcleanuphook()\n\ +\n\ +Return the the global cleanup hook." +); + +static PyObject * sys_setprofile(PyObject *self, PyObject *args) { if (trace_init() == -1) @@ -1092,6 +1141,8 @@ #endif {"settrace", sys_settrace, METH_O, settrace_doc}, {"gettrace", sys_gettrace, METH_NOARGS, gettrace_doc}, + {"setcleanuphook", sys_setcleanuphook, METH_VARARGS, setcleanuphook_doc}, + {"getcleanuphook", sys_getcleanuphook, METH_NOARGS, getcleanuphook_doc}, {"call_tracing", sys_call_tracing, METH_VARARGS, call_tracing_doc}, {NULL, NULL} /* sentinel */ };