1 /* 2 * atexit - allow programmer to define multiple exit functions to be executed 3 * upon normal program termination. 4 * 5 * Translated from atexit.py by Collin Winter. 6 + Copyright 2007 Python Software Foundation. 7 */ 8 9 #include "Python.h" 10 11 /* Forward declaration (for atexit_cleanup) */ 12 static PyObject *atexit_clear(PyObject*, PyObject*); 13 /* Forward declaration of module object */ 14 static struct PyModuleDef atexitmodule; 15 16 /* ===================================================================== */ 17 /* Callback machinery. */ 18 19 typedef struct { 20 PyObject *func; 21 PyObject *args; 22 PyObject *kwargs; 23 } atexit_callback; 24 25 typedef struct { 26 atexit_callback **atexit_callbacks; 27 int ncallbacks; 28 int callback_len; 29 } atexitmodule_state; 30 31 #define GET_ATEXIT_STATE(mod) ((atexitmodule_state*)PyModule_GetState(mod)) 32 33 34 static void 35 atexit_delete_cb(atexitmodule_state *modstate, int i) 36 { 37 atexit_callback *cb; 38 39 cb = modstate->atexit_callbacks[i]; 40 modstate->atexit_callbacks[i] = NULL; 41 Py_DECREF(cb->func); 42 Py_DECREF(cb->args); 43 Py_XDECREF(cb->kwargs); 44 PyMem_Free(cb); 45 } 46 47 /* Clear all callbacks without calling them */ 48 static void 49 atexit_cleanup(atexitmodule_state *modstate) 50 { 51 atexit_callback *cb; 52 int i; 53 for (i = 0; i < modstate->ncallbacks; i++) { 54 cb = modstate->atexit_callbacks[i]; 55 if (cb == NULL) 56 continue; 57 58 atexit_delete_cb(modstate, i); 59 } 60 modstate->ncallbacks = 0; 61 } 62 63 /* Installed into pylifecycle.c's atexit mechanism */ 64 65 static void 66 atexit_callfuncs(PyObject *module) 67 { 68 PyObject *exc_type = NULL, *exc_value, *exc_tb, *r; 69 atexit_callback *cb; 70 atexitmodule_state *modstate; 71 int i; 72 73 if (module == NULL) 74 return; 75 modstate = GET_ATEXIT_STATE(module); 76 77 if (modstate->ncallbacks == 0) 78 return; 79 80 81 for (i = modstate->ncallbacks - 1; i >= 0; i--) 82 { 83 cb = modstate->atexit_callbacks[i]; 84 if (cb == NULL) 85 continue; 86 87 r = PyObject_Call(cb->func, cb->args, cb->kwargs); 88 Py_XDECREF(r); 89 if (r == NULL) { 90 /* Maintain the last exception, but don't leak if there are 91 multiple exceptions. */ 92 if (exc_type) { 93 Py_DECREF(exc_type); 94 Py_XDECREF(exc_value); 95 Py_XDECREF(exc_tb); 96 } 97 PyErr_Fetch(&exc_type, &exc_value, &exc_tb); 98 if (!PyErr_GivenExceptionMatches(exc_type, PyExc_SystemExit)) { 99 PySys_WriteStderr("Error in atexit._run_exitfuncs:\n"); 100 PyErr_NormalizeException(&exc_type, &exc_value, &exc_tb); 101 PyErr_Display(exc_type, exc_value, exc_tb); 102 } 103 } 104 } 105 106 atexit_cleanup(modstate); 107 108 if (exc_type) 109 PyErr_Restore(exc_type, exc_value, exc_tb); 110 } 111 112 /* ===================================================================== */ 113 /* Module methods. */ 114 115 PyDoc_STRVAR(atexit_register__doc__, 116 "register(func, *args, **kwargs) -> func\n\ 117 \n\ 118 Register a function to be executed upon normal program termination\n\ 119 \n\ 120 func - function to be called at exit\n\ 121 args - optional arguments to pass to func\n\ 122 kwargs - optional keyword arguments to pass to func\n\ 123 \n\ 124 func is returned to facilitate usage as a decorator."); 125 126 static PyObject * 127 atexit_register(PyObject *self, PyObject *args, PyObject *kwargs) 128 { 129 atexitmodule_state *modstate; 130 atexit_callback *new_callback; 131 PyObject *func = NULL; 132 133 modstate = GET_ATEXIT_STATE(self); 134 135 if (modstate->ncallbacks >= modstate->callback_len) { 136 atexit_callback **r; 137 modstate->callback_len += 16; 138 r = (atexit_callback**)PyMem_Realloc(modstate->atexit_callbacks, 139 sizeof(atexit_callback*) * modstate->callback_len); 140 if (r == NULL) 141 return PyErr_NoMemory(); 142 modstate->atexit_callbacks = r; 143 } 144 145 if (PyTuple_GET_SIZE(args) == 0) { 146 PyErr_SetString(PyExc_TypeError, 147 "register() takes at least 1 argument (0 given)"); 148 return NULL; 149 } 150 151 func = PyTuple_GET_ITEM(args, 0); 152 if (!PyCallable_Check(func)) { 153 PyErr_SetString(PyExc_TypeError, 154 "the first argument must be callable"); 155 return NULL; 156 } 157 158 new_callback = PyMem_Malloc(sizeof(atexit_callback)); 159 if (new_callback == NULL) 160 return PyErr_NoMemory(); 161 162 new_callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args)); 163 if (new_callback->args == NULL) { 164 PyMem_Free(new_callback); 165 return NULL; 166 } 167 new_callback->func = func; 168 new_callback->kwargs = kwargs; 169 Py_INCREF(func); 170 Py_XINCREF(kwargs); 171 172 modstate->atexit_callbacks[modstate->ncallbacks++] = new_callback; 173 174 Py_INCREF(func); 175 return func; 176 } 177 178 PyDoc_STRVAR(atexit_run_exitfuncs__doc__, 179 "_run_exitfuncs() -> None\n\ 180 \n\ 181 Run all registered exit functions."); 182 183 static PyObject * 184 atexit_run_exitfuncs(PyObject *self, PyObject *unused) 185 { 186 atexit_callfuncs(self); 187 if (PyErr_Occurred()) 188 return NULL; 189 Py_RETURN_NONE; 190 } 191 192 PyDoc_STRVAR(atexit_clear__doc__, 193 "_clear() -> None\n\ 194 \n\ 195 Clear the list of previously registered exit functions."); 196 197 static PyObject * 198 atexit_clear(PyObject *self, PyObject *unused) 199 { 200 atexit_cleanup(GET_ATEXIT_STATE(self)); 201 Py_RETURN_NONE; 202 } 203 204 PyDoc_STRVAR(atexit_ncallbacks__doc__, 205 "_ncallbacks() -> int\n\ 206 \n\ 207 Return the number of registered exit functions."); 208 209 static PyObject * 210 atexit_ncallbacks(PyObject *self, PyObject *unused) 211 { 212 atexitmodule_state *modstate; 213 214 modstate = GET_ATEXIT_STATE(self); 215 216 return PyLong_FromSsize_t(modstate->ncallbacks); 217 } 218 219 static int 220 atexit_m_traverse(PyObject *self, visitproc visit, void *arg) 221 { 222 int i; 223 atexitmodule_state *modstate; 224 225 modstate = GET_ATEXIT_STATE(self); 226 if (modstate != NULL) { 227 for (i = 0; i < modstate->ncallbacks; i++) { 228 atexit_callback *cb = modstate->atexit_callbacks[i]; 229 if (cb == NULL) 230 continue; 231 Py_VISIT(cb->func); 232 Py_VISIT(cb->args); 233 Py_VISIT(cb->kwargs); 234 } 235 } 236 return 0; 237 } 238 239 static int 240 atexit_m_clear(PyObject *self) 241 { 242 atexitmodule_state *modstate; 243 modstate = GET_ATEXIT_STATE(self); 244 if (modstate != NULL) { 245 atexit_cleanup(modstate); 246 } 247 return 0; 248 } 249 250 static void 251 atexit_free(PyObject *m) 252 { 253 atexitmodule_state *modstate; 254 modstate = GET_ATEXIT_STATE(m); 255 if (modstate != NULL) { 256 atexit_cleanup(modstate); 257 PyMem_Free(modstate->atexit_callbacks); 258 } 259 } 260 261 PyDoc_STRVAR(atexit_unregister__doc__, 262 "unregister(func) -> None\n\ 263 \n\ 264 Unregister an exit function which was previously registered using\n\ 265 atexit.register\n\ 266 \n\ 267 func - function to be unregistered"); 268 269 static PyObject * 270 atexit_unregister(PyObject *self, PyObject *func) 271 { 272 atexitmodule_state *modstate; 273 atexit_callback *cb; 274 int i, eq; 275 276 modstate = GET_ATEXIT_STATE(self); 277 278 for (i = 0; i < modstate->ncallbacks; i++) 279 { 280 cb = modstate->atexit_callbacks[i]; 281 if (cb == NULL) 282 continue; 283 284 eq = PyObject_RichCompareBool(cb->func, func, Py_EQ); 285 if (eq < 0) 286 return NULL; 287 if (eq) 288 atexit_delete_cb(modstate, i); 289 } 290 Py_RETURN_NONE; 291 } 292 293 static PyMethodDef atexit_methods[] = { 294 {"register", (PyCFunction) atexit_register, METH_VARARGS|METH_KEYWORDS, 295 atexit_register__doc__}, 296 {"_clear", (PyCFunction) atexit_clear, METH_NOARGS, 297 atexit_clear__doc__}, 298 {"unregister", (PyCFunction) atexit_unregister, METH_O, 299 atexit_unregister__doc__}, 300 {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS, 301 atexit_run_exitfuncs__doc__}, 302 {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS, 303 atexit_ncallbacks__doc__}, 304 {NULL, NULL} /* sentinel */ 305 }; 306 307 /* ===================================================================== */ 308 /* Initialization function. */ 309 310 PyDoc_STRVAR(atexit__doc__, 311 "allow programmer to define multiple exit functions to be executed\ 312 upon normal program termination.\n\ 313 \n\ 314 Two public functions, register and unregister, are defined.\n\ 315 "); 316 317 static int 318 atexit_exec(PyObject *m) { 319 atexitmodule_state *modstate; 320 321 modstate = GET_ATEXIT_STATE(m); 322 modstate->callback_len = 32; 323 modstate->ncallbacks = 0; 324 modstate->atexit_callbacks = PyMem_New(atexit_callback*, 325 modstate->callback_len); 326 if (modstate->atexit_callbacks == NULL) 327 return -1; 328 329 _Py_PyAtExit(atexit_callfuncs, m); 330 return 0; 331 } 332 333 static PyModuleDef_Slot atexit_slots[] = { 334 {Py_mod_exec, atexit_exec}, 335 {0, NULL} 336 }; 337 338 static struct PyModuleDef atexitmodule = { 339 PyModuleDef_HEAD_INIT, 340 "atexit", 341 atexit__doc__, 342 sizeof(atexitmodule_state), 343 atexit_methods, 344 atexit_slots, 345 atexit_m_traverse, 346 atexit_m_clear, 347 (freefunc)atexit_free 348 }; 349 350 PyMODINIT_FUNC 351 PyInit_atexit(void) 352 { 353 return PyModuleDef_Init(&atexitmodule); 354 } 355