Home | History | Annotate | Download | only in Modules
      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