Modules/atexitmodule.c

   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*);
  13
  14 /* ===================================================================== */
  15 /* Callback machinery. */
  16
  17 typedef struct {
  18     PyObject *func;
  19     PyObject *args;
  20     PyObject *kwargs;
  21 } atexit_callback;
  22
  23 static atexit_callback **atexit_callbacks;
  24 static int ncallbacks = 0;
  25 static int callback_len = 32;
  26
  27 /* Installed into pythonrun.c's atexit mechanism */
  28
  29 void
  30 atexit_callfuncs(void)
  31 {
  32     PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;
  33     atexit_callback *cb;
  34     int i;
  35
  36     if (ncallbacks == 0)
  37         return;
  38
  39     for (i = ncallbacks - 1; i >= 0; i--)
  40     {
  41         cb = atexit_callbacks[i];
  42         if (cb == NULL)
  43             continue;
  44
  45         r = PyObject_Call(cb->func, cb->args, cb->kwargs);
  46         Py_XDECREF(r);
  47         if (r == NULL) {
  48             /* Maintain the last exception, but don't leak if there are
  49                multiple exceptions. */
  50             if (exc_type) {
  51                 Py_DECREF(exc_type);
  52                 Py_XDECREF(exc_value);
  53                 Py_XDECREF(exc_tb);
  54             }
  55             PyErr_Fetch(&exc_type, &exc_value, &exc_tb);
  56             if (!PyErr_ExceptionMatches(PyExc_SystemExit)) {
  57                 PySys_WriteStderr("Error in atexit._run_exitfuncs:\n");
  58                 PyErr_Display(exc_type, exc_value, exc_tb);
  59             }
  60         }
  61     }
  62
  63     if (exc_type)
  64         PyErr_Restore(exc_type, exc_value, exc_tb);
  65 }
  66
  67 void
  68 atexit_delete_cb(int i)
  69 {
  70     atexit_callback *cb = atexit_callbacks[i];
  71     atexit_callbacks[i] = NULL;
  72     Py_DECREF(cb->func);
  73     Py_DECREF(cb->args);
  74     Py_XDECREF(cb->kwargs);
  75     PyMem_Free(cb);
  76 }
  77
  78 void
  79 atexit_cleanup(void)
  80 {
  81     PyObject *r = atexit_clear(NULL);
  82     Py_DECREF(r);
  83 }
  84
  85 /* ===================================================================== */
  86 /* Module methods. */
  87
  88 PyDoc_STRVAR(atexit_register__doc__,
  89 "register(func, *args, **kwargs) -> func\n\
  90 \n\
  91 Register a function to be executed upon normal program termination\n\
  92 \n\
  93     func - function to be called at exit\n\
  94     args - optional arguments to pass to func\n\
  95     kwargs - optional keyword arguments to pass to func\n\
  96 \n\
  97     func is returned to facilitate usage as a decorator.");
  98
  99 static PyObject *
 100 atexit_register(PyObject *self, PyObject *args, PyObject *kwargs)
 101 {
 102     atexit_callback *new_callback;
 103     PyObject *func = NULL;
 104
 105     if (ncallbacks >= callback_len) {
 106         atexit_callback **r;
 107         callback_len += 16;
 108         r = (atexit_callback**)PyMem_Realloc(atexit_callbacks,
 109                                       sizeof(atexit_callback*) * callback_len);
 110         if (r == NULL)
 111             return PyErr_NoMemory();
 112         atexit_callbacks = r;
 113     }
 114
 115     if (PyTuple_GET_SIZE(args) == 0) {
 116         PyErr_SetString(PyExc_TypeError,
 117                 "register() takes at least 1 argument (0 given)");
 118         return NULL;
 119     }
 120
 121     func = PyTuple_GET_ITEM(args, 0);
 122     if (!PyCallable_Check(func)) {
 123         PyErr_SetString(PyExc_TypeError,
 124                 "the first argument must be callable");
 125         return NULL;
 126     }
 127
 128     new_callback = PyMem_Malloc(sizeof(atexit_callback));
 129     if (new_callback == NULL)
 130         return PyErr_NoMemory();
 131
 132     new_callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
 133     if (new_callback->args == NULL) {
 134         PyMem_Free(new_callback);
 135         return NULL;
 136     }
 137     new_callback->func = func;
 138     new_callback->kwargs = kwargs;
 139     Py_INCREF(func);
 140     Py_XINCREF(kwargs);
 141
 142     atexit_callbacks[ncallbacks++] = new_callback;
 143
 144     Py_INCREF(func);
 145     return func;
 146 }
 147
 148 PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
 149 "_run_exitfuncs() -> None\n\
 150 \n\
 151 Run all registered exit functions.");
 152
 153 static PyObject *
 154 atexit_run_exitfuncs(PyObject *self)
 155 {
 156     atexit_callfuncs();
 157     if (PyErr_Occurred())
 158         return NULL;
 159     Py_RETURN_NONE;
 160 }
 161
 162 PyDoc_STRVAR(atexit_clear__doc__,
 163 "_clear() -> None\n\
 164 \n\
 165 Clear the list of previously registered exit functions.");
 166
 167 static PyObject *
 168 atexit_clear(PyObject *self)
 169 {
 170     atexit_callback *cb;
 171     int i;
 172
 173     for (i = 0; i < ncallbacks; i++)
 174     {
 175         cb = atexit_callbacks[i];
 176         if (cb == NULL)
 177             continue;
 178
 179         atexit_delete_cb(i);
 180     }
 181     ncallbacks = 0;
 182     Py_RETURN_NONE;
 183 }
 184
 185 PyDoc_STRVAR(atexit_unregister__doc__,
 186 "unregister(func) -> None\n\
 187 \n\
 188 Unregister a exit function which was previously registered using\n\
 189 atexit.register\n\
 190 \n\
 191     func - function to be unregistered");
 192
 193 static PyObject *
 194 atexit_unregister(PyObject *self, PyObject *func)
 195 {
 196     atexit_callback *cb;
 197     int i, eq;
 198
 199     for (i = 0; i < ncallbacks; i++)
 200     {
 201         cb = atexit_callbacks[i];
 202         if (cb == NULL)
 203             continue;
 204
 205         eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
 206         if (eq < 0)
 207             return NULL;
 208         if (eq)
 209             atexit_delete_cb(i);
 210     }
 211     Py_RETURN_NONE;
 212 }
 213
 214 static PyMethodDef atexit_methods[] = {
 215     {"register", (PyCFunction) atexit_register, METH_VARARGS|METH_KEYWORDS,
 216         atexit_register__doc__},
 217     {"_clear", (PyCFunction) atexit_clear, METH_NOARGS,
 218         atexit_clear__doc__},
 219     {"unregister", (PyCFunction) atexit_unregister, METH_O,
 220         atexit_unregister__doc__},
 221     {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS,
 222         atexit_run_exitfuncs__doc__},
 223     {NULL, NULL}        /* sentinel */
 224 };
 225
 226 /* ===================================================================== */
 227 /* Initialization function. */
 228
 229 PyDoc_STRVAR(atexit__doc__,
 230 "allow programmer to define multiple exit functions to be executed\
 231 upon normal program termination.\n\
 232 \n\
 233 Two public functions, register and unregister, are defined.\n\
 234 ");
 235
 236 PyMODINIT_FUNC
 237 initatexit(void)
 238 {
 239     PyObject *m;
 240
 241     atexit_callbacks = PyMem_New(atexit_callback*, callback_len);
 242     if (atexit_callbacks == NULL)
 243         return;
 244
 245     m = Py_InitModule3("atexit", atexit_methods, atexit__doc__);
 246     if (m == NULL)
 247         return;
 248
 249     _Py_PyAtExit(atexit_callfuncs);
 250     /* Register a callback that will free
 251        atexit_callbacks, otherwise valgrind will report memory leaks. */
 252     Py_AtExit(atexit_cleanup);
 253 }