/*

 *  atexit - allow programmer to define multiple exit functions to be executed

 *  upon normal program termination.

 *

 *   Translated from atexit.py by Collin Winter.

 +   Copyright 2007 Python Software Foundation.

 */

#include "Python.h"

/* Forward declaration (for atexit_cleanup) */

static PyObject *atexit_clear(PyObject*, PyObject*);

/* Forward declaration of module object */

static struct PyModuleDef atexitmodule;

/* ===================================================================== */

/* Callback machinery. */

typedef struct {

    PyObject *func;

    PyObject *args;

    PyObject *kwargs;

} atexit_callback;

typedef struct {

    atexit_callback **atexit_callbacks;

    int ncallbacks;

    int callback_len;

} atexitmodule_state;

#define GET_ATEXIT_STATE(mod) ((atexitmodule_state*)PyModule_GetState(mod))

static void

atexit_delete_cb(atexitmodule_state *modstate, int i)

{

    atexit_callback *cb;

    cb = modstate->atexit_callbacks[i];

    modstate->atexit_callbacks[i] = NULL;

    Py_DECREF(cb->func);

    Py_DECREF(cb->args);

    Py_XDECREF(cb->kwargs);

    PyMem_Free(cb);

}

/* Clear all callbacks without calling them */

static void

atexit_cleanup(atexitmodule_state *modstate)

{

    atexit_callback *cb;

    int i;

    for (i = 0; i < modstate->ncallbacks; i++) {

        cb = modstate->atexit_callbacks[i];

        if (cb == NULL)

            continue;

        atexit_delete_cb(modstate, i);

    }

    modstate->ncallbacks = 0;

}

/* Installed into pylifecycle.c's atexit mechanism */

static void

atexit_callfuncs(PyObject *module)

{

    PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;

    atexit_callback *cb;

    atexitmodule_state *modstate;

    int i;

    if (module == NULL)

        return;

    modstate = GET_ATEXIT_STATE(module);

    if (modstate->ncallbacks == 0)

        return;

    for (i = modstate->ncallbacks - 1; i >= 0; i--)

    {

        cb = modstate->atexit_callbacks[i];

        if (cb == NULL)

            continue;

        r = PyObject_Call(cb->func, cb->args, cb->kwargs);

        Py_XDECREF(r);

        if (r == NULL) {

            /* Maintain the last exception, but don't leak if there are

               multiple exceptions. */

            if (exc_type) {

                Py_DECREF(exc_type);

                Py_XDECREF(exc_value);

                Py_XDECREF(exc_tb);

            }

            PyErr_Fetch(&exc_type, &exc_value, &exc_tb);

            if (!PyErr_GivenExceptionMatches(exc_type, PyExc_SystemExit)) {

                PySys_WriteStderr("Error in atexit._run_exitfuncs:\n");

                PyErr_NormalizeException(&exc_type, &exc_value, &exc_tb);

                PyErr_Display(exc_type, exc_value, exc_tb);

            }

        }

    }

    atexit_cleanup(modstate);

    if (exc_type)

        PyErr_Restore(exc_type, exc_value, exc_tb);

}

/* ===================================================================== */

/* Module methods. */

PyDoc_STRVAR(atexit_register__doc__,

"register(func, *args, **kwargs) -> func\n\

\n\

Register a function to be executed upon normal program termination\n\

\n\

    func - function to be called at exit\n\

    args - optional arguments to pass to func\n\

    kwargs - optional keyword arguments to pass to func\n\

\n\

    func is returned to facilitate usage as a decorator.");

static PyObject *

atexit_register(PyObject *self, PyObject *args, PyObject *kwargs)

{

    atexitmodule_state *modstate;

    atexit_callback *new_callback;

    PyObject *func = NULL;

    modstate = GET_ATEXIT_STATE(self);

    if (modstate->ncallbacks >= modstate->callback_len) {

        atexit_callback **r;

        modstate->callback_len += 16;

        r = (atexit_callback**)PyMem_Realloc(modstate->atexit_callbacks,

                                      sizeof(atexit_callback*) * modstate->callback_len);

        if (r == NULL)

            return PyErr_NoMemory();

        modstate->atexit_callbacks = r;

    }

    if (PyTuple_GET_SIZE(args) == 0) {

        PyErr_SetString(PyExc_TypeError,

                "register() takes at least 1 argument (0 given)");

        return NULL;

    }

    func = PyTuple_GET_ITEM(args, 0);

    if (!PyCallable_Check(func)) {

        PyErr_SetString(PyExc_TypeError,

                "the first argument must be callable");

        return NULL;

    }

    new_callback = PyMem_Malloc(sizeof(atexit_callback));

    if (new_callback == NULL)

        return PyErr_NoMemory();

    new_callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));

    if (new_callback->args == NULL) {

        PyMem_Free(new_callback);

        return NULL;

    }

    new_callback->func = func;

    new_callback->kwargs = kwargs;

    Py_INCREF(func);

    Py_XINCREF(kwargs);

    modstate->atexit_callbacks[modstate->ncallbacks++] = new_callback;

    Py_INCREF(func);

    return func;

}

PyDoc_STRVAR(atexit_run_exitfuncs__doc__,

"_run_exitfuncs() -> None\n\

\n\

Run all registered exit functions.");

static PyObject *

atexit_run_exitfuncs(PyObject *self, PyObject *unused)

{

    atexit_callfuncs(self);

    if (PyErr_Occurred())

        return NULL;

    Py_RETURN_NONE;

}

PyDoc_STRVAR(atexit_clear__doc__,

"_clear() -> None\n\

\n\

Clear the list of previously registered exit functions.");

static PyObject *

atexit_clear(PyObject *self, PyObject *unused)

{

    atexit_cleanup(GET_ATEXIT_STATE(self));

    Py_RETURN_NONE;

}

PyDoc_STRVAR(atexit_ncallbacks__doc__,

"_ncallbacks() -> int\n\

\n\

Return the number of registered exit functions.");

static PyObject *

atexit_ncallbacks(PyObject *self, PyObject *unused)

{

    atexitmodule_state *modstate;

    modstate = GET_ATEXIT_STATE(self);

    return PyLong_FromSsize_t(modstate->ncallbacks);

}

static int

atexit_m_traverse(PyObject *self, visitproc visit, void *arg)

{

    int i;

    atexitmodule_state *modstate;

    modstate = GET_ATEXIT_STATE(self);

    if (modstate != NULL) {

        for (i = 0; i < modstate->ncallbacks; i++) {

            atexit_callback *cb = modstate->atexit_callbacks[i];

            if (cb == NULL)

                continue;

            Py_VISIT(cb->func);

            Py_VISIT(cb->args);

            Py_VISIT(cb->kwargs);

        }

    }

    return 0;

}

static int

atexit_m_clear(PyObject *self)

{

    atexitmodule_state *modstate;

    modstate = GET_ATEXIT_STATE(self);

    if (modstate != NULL) {

        atexit_cleanup(modstate);

    }

    return 0;

}

static void

atexit_free(PyObject *m)

{

    atexitmodule_state *modstate;

    modstate = GET_ATEXIT_STATE(m);

    if (modstate != NULL) {

        atexit_cleanup(modstate);

        PyMem_Free(modstate->atexit_callbacks);

    }

}

PyDoc_STRVAR(atexit_unregister__doc__,

"unregister(func) -> None\n\

\n\

Unregister an exit function which was previously registered using\n\

atexit.register\n\

\n\

    func - function to be unregistered");

static PyObject *

atexit_unregister(PyObject *self, PyObject *func)

{

    atexitmodule_state *modstate;

    atexit_callback *cb;

    int i, eq;

    modstate = GET_ATEXIT_STATE(self);

    for (i = 0; i < modstate->ncallbacks; i++)

    {

        cb = modstate->atexit_callbacks[i];

        if (cb == NULL)

            continue;

        eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);

        if (eq < 0)

            return NULL;

        if (eq)

            atexit_delete_cb(modstate, i);

    }

    Py_RETURN_NONE;

}

static PyMethodDef atexit_methods[] = {

    {"register", (PyCFunction)(void(*)(void)) atexit_register, METH_VARARGS|METH_KEYWORDS,

        atexit_register__doc__},

    {"_clear", (PyCFunction) atexit_clear, METH_NOARGS,

        atexit_clear__doc__},

    {"unregister", (PyCFunction) atexit_unregister, METH_O,

        atexit_unregister__doc__},

    {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS,

        atexit_run_exitfuncs__doc__},

    {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS,

        atexit_ncallbacks__doc__},

    {NULL, NULL}        /* sentinel */

};

/* ===================================================================== */

/* Initialization function. */

PyDoc_STRVAR(atexit__doc__,

"allow programmer to define multiple exit functions to be executed\

upon normal program termination.\n\

\n\

Two public functions, register and unregister, are defined.\n\

");

static int

atexit_exec(PyObject *m) {

    atexitmodule_state *modstate;

    modstate = GET_ATEXIT_STATE(m);

    modstate->callback_len = 32;

    modstate->ncallbacks = 0;

    modstate->atexit_callbacks = PyMem_New(atexit_callback*,

                                           modstate->callback_len);

    if (modstate->atexit_callbacks == NULL)

        return -1;

    _Py_PyAtExit(atexit_callfuncs, m);

    return 0;

}

static PyModuleDef_Slot atexit_slots[] = {

    {Py_mod_exec, atexit_exec},

    {0, NULL}

};

static struct PyModuleDef atexitmodule = {

    PyModuleDef_HEAD_INIT,

    "atexit",

    atexit__doc__,

    sizeof(atexitmodule_state),

    atexit_methods,

    atexit_slots,

    atexit_m_traverse,

    atexit_m_clear,

    (freefunc)atexit_free

};

PyMODINIT_FUNC

PyInit_atexit(void)

{

    return PyModuleDef_Init(&atexitmodule);

}