Line data    Source code

       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             : #include "pycore_initconfig.h"    // _PyStatus_NO_MEMORY
      11             : #include "pycore_interp.h"        // PyInterpreterState.atexit
      12             : #include "pycore_pystate.h"       // _PyInterpreterState_GET
      13             : 
      14             : /* ===================================================================== */
      15             : /* Callback machinery. */
      16             : 
      17             : static inline struct atexit_state*
      18        2639 : get_atexit_state(void)
      19             : {
      20        2639 :     PyInterpreterState *interp = _PyInterpreterState_GET();
      21        2639 :     return &interp->atexit;
      22             : }
      23             : 
      24             : 
      25             : static void
      26        2269 : atexit_delete_cb(struct atexit_state *state, int i)
      27             : {
      28        2269 :     atexit_callback *cb = state->callbacks[i];
      29        2269 :     state->callbacks[i] = NULL;
      30             : 
      31        2269 :     Py_DECREF(cb->func);
      32        2269 :     Py_DECREF(cb->args);
      33        2269 :     Py_XDECREF(cb->kwargs);
      34        2269 :     PyMem_Free(cb);
      35        2269 : }
      36             : 
      37             : 
      38             : /* Clear all callbacks without calling them */
      39             : static void
      40        4346 : atexit_cleanup(struct atexit_state *state)
      41             : {
      42             :     atexit_callback *cb;
      43        6615 :     for (int i = 0; i < state->ncallbacks; i++) {
      44        2269 :         cb = state->callbacks[i];
      45        2269 :         if (cb == NULL)
      46         335 :             continue;
      47             : 
      48        1934 :         atexit_delete_cb(state, i);
      49             :     }
      50        4346 :     state->ncallbacks = 0;
      51        4346 : }
      52             : 
      53             : 
      54             : PyStatus
      55        3134 : _PyAtExit_Init(PyInterpreterState *interp)
      56             : {
      57        3134 :     struct atexit_state *state = &interp->atexit;
      58             :     // _PyAtExit_Init() must only be called once
      59        3134 :     assert(state->callbacks == NULL);
      60             : 
      61        3134 :     state->callback_len = 32;
      62        3134 :     state->ncallbacks = 0;
      63        3134 :     state->callbacks = PyMem_New(atexit_callback*, state->callback_len);
      64        3134 :     if (state->callbacks == NULL) {
      65           0 :         return _PyStatus_NO_MEMORY();
      66             :     }
      67        3134 :     return _PyStatus_OK();
      68             : }
      69             : 
      70             : 
      71             : void
      72        3120 : _PyAtExit_Fini(PyInterpreterState *interp)
      73             : {
      74        3120 :     struct atexit_state *state = &interp->atexit;
      75        3120 :     atexit_cleanup(state);
      76        3120 :     PyMem_Free(state->callbacks);
      77        3120 :     state->callbacks = NULL;
      78        3120 : }
      79             : 
      80             : 
      81             : static void
      82        3131 : atexit_callfuncs(struct atexit_state *state)
      83             : {
      84        3131 :     assert(!PyErr_Occurred());
      85             : 
      86        3131 :     if (state->ncallbacks == 0) {
      87        1926 :         return;
      88             :     }
      89             : 
      90        3473 :     for (int i = state->ncallbacks - 1; i >= 0; i--) {
      91        2268 :         atexit_callback *cb = state->callbacks[i];
      92        2268 :         if (cb == NULL) {
      93         334 :             continue;
      94             :         }
      95             : 
      96             :         // bpo-46025: Increment the refcount of cb->func as the call itself may unregister it
      97        1934 :         PyObject* the_func = Py_NewRef(cb->func);
      98        1934 :         PyObject *res = PyObject_Call(cb->func, cb->args, cb->kwargs);
      99        1934 :         if (res == NULL) {
     100          10 :             _PyErr_WriteUnraisableMsg("in atexit callback", the_func);
     101             :         }
     102             :         else {
     103        1924 :             Py_DECREF(res);
     104             :         }
     105        1934 :         Py_DECREF(the_func);
     106             :     }
     107             : 
     108        1205 :     atexit_cleanup(state);
     109             : 
     110        1205 :     assert(!PyErr_Occurred());
     111             : }
     112             : 
     113             : 
     114             : void
     115        3120 : _PyAtExit_Call(PyInterpreterState *interp)
     116             : {
     117        3120 :     struct atexit_state *state = &interp->atexit;
     118        3120 :     atexit_callfuncs(state);
     119        3120 : }
     120             : 
     121             : 
     122             : /* ===================================================================== */
     123             : /* Module methods. */
     124             : 
     125             : 
     126             : PyDoc_STRVAR(atexit_register__doc__,
     127             : "register(func, *args, **kwargs) -> func\n\
     128             : \n\
     129             : Register a function to be executed upon normal program termination\n\
     130             : \n\
     131             :     func - function to be called at exit\n\
     132             :     args - optional arguments to pass to func\n\
     133             :     kwargs - optional keyword arguments to pass to func\n\
     134             : \n\
     135             :     func is returned to facilitate usage as a decorator.");
     136             : 
     137             : static PyObject *
     138        2269 : atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)
     139             : {
     140        2269 :     if (PyTuple_GET_SIZE(args) == 0) {
     141           0 :         PyErr_SetString(PyExc_TypeError,
     142             :                 "register() takes at least 1 argument (0 given)");
     143           0 :         return NULL;
     144             :     }
     145             : 
     146        2269 :     PyObject *func = PyTuple_GET_ITEM(args, 0);
     147        2269 :     if (!PyCallable_Check(func)) {
     148           0 :         PyErr_SetString(PyExc_TypeError,
     149             :                 "the first argument must be callable");
     150           0 :         return NULL;
     151             :     }
     152             : 
     153        2269 :     struct atexit_state *state = get_atexit_state();
     154        2269 :     if (state->ncallbacks >= state->callback_len) {
     155             :         atexit_callback **r;
     156           6 :         state->callback_len += 16;
     157           6 :         size_t size = sizeof(atexit_callback*) * (size_t)state->callback_len;
     158           6 :         r = (atexit_callback**)PyMem_Realloc(state->callbacks, size);
     159           6 :         if (r == NULL) {
     160           0 :             return PyErr_NoMemory();
     161             :         }
     162           6 :         state->callbacks = r;
     163             :     }
     164             : 
     165        2269 :     atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
     166        2269 :     if (callback == NULL) {
     167           0 :         return PyErr_NoMemory();
     168             :     }
     169             : 
     170        2269 :     callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
     171        2269 :     if (callback->args == NULL) {
     172           0 :         PyMem_Free(callback);
     173           0 :         return NULL;
     174             :     }
     175        2269 :     callback->func = Py_NewRef(func);
     176        2269 :     callback->kwargs = Py_XNewRef(kwargs);
     177             : 
     178        2269 :     state->callbacks[state->ncallbacks++] = callback;
     179             : 
     180        2269 :     return Py_NewRef(func);
     181             : }
     182             : 
     183             : PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
     184             : "_run_exitfuncs() -> None\n\
     185             : \n\
     186             : Run all registered exit functions.\n\
     187             : \n\
     188             : If a callaback raises an exception, it is logged with sys.unraisablehook.");
     189             : 
     190             : static PyObject *
     191          11 : atexit_run_exitfuncs(PyObject *module, PyObject *unused)
     192             : {
     193          11 :     struct atexit_state *state = get_atexit_state();
     194          11 :     atexit_callfuncs(state);
     195          11 :     Py_RETURN_NONE;
     196             : }
     197             : 
     198             : PyDoc_STRVAR(atexit_clear__doc__,
     199             : "_clear() -> None\n\
     200             : \n\
     201             : Clear the list of previously registered exit functions.");
     202             : 
     203             : static PyObject *
     204          21 : atexit_clear(PyObject *module, PyObject *unused)
     205             : {
     206          21 :     atexit_cleanup(get_atexit_state());
     207          21 :     Py_RETURN_NONE;
     208             : }
     209             : 
     210             : PyDoc_STRVAR(atexit_ncallbacks__doc__,
     211             : "_ncallbacks() -> int\n\
     212             : \n\
     213             : Return the number of registered exit functions.");
     214             : 
     215             : static PyObject *
     216           4 : atexit_ncallbacks(PyObject *module, PyObject *unused)
     217             : {
     218           4 :     struct atexit_state *state = get_atexit_state();
     219           4 :     return PyLong_FromSsize_t(state->ncallbacks);
     220             : }
     221             : 
     222             : PyDoc_STRVAR(atexit_unregister__doc__,
     223             : "unregister(func) -> None\n\
     224             : \n\
     225             : Unregister an exit function which was previously registered using\n\
     226             : atexit.register\n\
     227             : \n\
     228             :     func - function to be unregistered");
     229             : 
     230             : static PyObject *
     231         334 : atexit_unregister(PyObject *module, PyObject *func)
     232             : {
     233         334 :     struct atexit_state *state = get_atexit_state();
     234        1003 :     for (int i = 0; i < state->ncallbacks; i++)
     235             :     {
     236         669 :         atexit_callback *cb = state->callbacks[i];
     237         669 :         if (cb == NULL) {
     238           0 :             continue;
     239             :         }
     240             : 
     241         669 :         int eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
     242         669 :         if (eq < 0) {
     243           0 :             return NULL;
     244             :         }
     245         669 :         if (eq) {
     246         335 :             atexit_delete_cb(state, i);
     247             :         }
     248             :     }
     249         334 :     Py_RETURN_NONE;
     250             : }
     251             : 
     252             : 
     253             : static PyMethodDef atexit_methods[] = {
     254             :     {"register", _PyCFunction_CAST(atexit_register), METH_VARARGS|METH_KEYWORDS,
     255             :         atexit_register__doc__},
     256             :     {"_clear", (PyCFunction) atexit_clear, METH_NOARGS,
     257             :         atexit_clear__doc__},
     258             :     {"unregister", (PyCFunction) atexit_unregister, METH_O,
     259             :         atexit_unregister__doc__},
     260             :     {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS,
     261             :         atexit_run_exitfuncs__doc__},
     262             :     {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS,
     263             :         atexit_ncallbacks__doc__},
     264             :     {NULL, NULL}        /* sentinel */
     265             : };
     266             : 
     267             : 
     268             : /* ===================================================================== */
     269             : /* Initialization function. */
     270             : 
     271             : PyDoc_STRVAR(atexit__doc__,
     272             : "allow programmer to define multiple exit functions to be executed\n\
     273             : upon normal program termination.\n\
     274             : \n\
     275             : Two public functions, register and unregister, are defined.\n\
     276             : ");
     277             : 
     278             : static struct PyModuleDef atexitmodule = {
     279             :     PyModuleDef_HEAD_INIT,
     280             :     .m_name = "atexit",
     281             :     .m_doc = atexit__doc__,
     282             :     .m_size = 0,
     283             :     .m_methods = atexit_methods,
     284             : };
     285             : 
     286             : PyMODINIT_FUNC
     287        1199 : PyInit_atexit(void)
     288             : {
     289        1199 :     return PyModuleDef_Init(&atexitmodule);
     290             : }