bpo-42639: Cleanup atexitmodule.c (GH-23770)

* Rename "atexitmodule_state" to "struct atexit_state".
* Rename "modstate" to "state".
* Rename "self" parameter to "module".
* test_atexit uses textwrap.dedent().
* Remove _Py_PyAtExit() function: inline it into atexit_exec().
* PyInterpreterState: rename pyexitfunc to atexit_func, rename
  pyexitmodule to atexit_module.

1 files changed, 109 insertions, 126 deletions

diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c
index fd47ddd..f1fc871 100644
--- a/Modules/atexitmodule.c
+++ b/Modules/atexitmodule.c
@@ -7,11 +7,8 @@
  */
 
 #include "Python.h"
-
-/* Forward declaration (for atexit_cleanup) */
-static PyObject *atexit_clear(PyObject*, PyObject*);
-/* Forward declaration of module object */
-static struct PyModuleDef atexitmodule;
+#include "pycore_interp.h"        // PyInterpreterState.atexit_func
+#include "pycore_pystate.h"       // _PyInterpreterState_GET
 
 /* ===================================================================== */
 /* Callback machinery. */
@@ -22,48 +19,47 @@ typedef struct {
     PyObject *kwargs;
 } atexit_callback;
 
-typedef struct {
+struct atexit_state {
     atexit_callback **atexit_callbacks;
     int ncallbacks;
     int callback_len;
-} atexitmodule_state;
+};
 
-static inline atexitmodule_state*
+static inline struct atexit_state*
 get_atexit_state(PyObject *module)
 {
     void *state = PyModule_GetState(module);
     assert(state != NULL);
-    return (atexitmodule_state *)state;
+    return (struct atexit_state *)state;
 }
 
 
 static void
-atexit_delete_cb(atexitmodule_state *modstate, int i)
+atexit_delete_cb(struct atexit_state *state, int i)
 {
-    atexit_callback *cb;
+    atexit_callback *cb = state->atexit_callbacks[i];
+    state->atexit_callbacks[i] = NULL;
 
-    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_cleanup(struct atexit_state *state)
 {
     atexit_callback *cb;
-    int i;
-    for (i = 0; i < modstate->ncallbacks; i++) {
-        cb = modstate->atexit_callbacks[i];
+    for (int i = 0; i < state->ncallbacks; i++) {
+        cb = state->atexit_callbacks[i];
         if (cb == NULL)
             continue;
 
-        atexit_delete_cb(modstate, i);
+        atexit_delete_cb(state, i);
     }
-    modstate->ncallbacks = 0;
+    state->ncallbacks = 0;
 }
 
 /* Installed into pylifecycle.c's atexit mechanism */
@@ -71,28 +67,26 @@ atexit_cleanup(atexitmodule_state *modstate)
 static void
 atexit_callfuncs(PyObject *module)
 {
-    PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;
-    atexit_callback *cb;
-    atexitmodule_state *modstate;
-    int i;
+    assert(!PyErr_Occurred());
 
-    if (module == NULL)
+    if (module == NULL) {
         return;
-    modstate = get_atexit_state(module);
+    }
 
-    if (modstate->ncallbacks == 0)
+    struct atexit_state *state = get_atexit_state(module);
+    if (state->ncallbacks == 0) {
         return;
+    }
 
-
-    for (i = modstate->ncallbacks - 1; i >= 0; i--)
-    {
-        cb = modstate->atexit_callbacks[i];
-        if (cb == NULL)
+    PyObject *exc_type = NULL, *exc_value, *exc_tb;
+    for (int i = state->ncallbacks - 1; i >= 0; i--) {
+        atexit_callback *cb = state->atexit_callbacks[i];
+        if (cb == NULL) {
             continue;
+        }
 
-        r = PyObject_Call(cb->func, cb->args, cb->kwargs);
-        Py_XDECREF(r);
-        if (r == NULL) {
+        PyObject *res = PyObject_Call(cb->func, cb->args, cb->kwargs);
+        if (res == NULL) {
             /* Maintain the last exception, but don't leak if there are
                multiple exceptions. */
             if (exc_type) {
@@ -107,12 +101,16 @@ atexit_callfuncs(PyObject *module)
                 PyErr_Display(exc_type, exc_value, exc_tb);
             }
         }
+        else {
+            Py_DECREF(res);
+        }
     }
 
-    atexit_cleanup(modstate);
+    atexit_cleanup(state);
 
-    if (exc_type)
+    if (exc_type) {
         PyErr_Restore(exc_type, exc_value, exc_tb);
+    }
 }
 
 /* ===================================================================== */
@@ -130,55 +128,48 @@ Register a function to be executed upon normal program termination\n\
     func is returned to facilitate usage as a decorator.");
 
 static PyObject *
-atexit_register(PyObject *self, PyObject *args, PyObject *kwargs)
+atexit_register(PyObject *module, 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);
+    PyObject *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)
+    struct atexit_state *state = get_atexit_state(module);
+    if (state->ncallbacks >= state->callback_len) {
+        atexit_callback **r;
+        state->callback_len += 16;
+        r = (atexit_callback**)PyMem_Realloc(state->atexit_callbacks,
+                                      sizeof(atexit_callback*) * state->callback_len);
+        if (r == NULL)
+            return PyErr_NoMemory();
+        state->atexit_callbacks = r;
+    }
+
+    atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
+    if (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);
+    callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
+    if (callback->args == NULL) {
+        PyMem_Free(callback);
         return NULL;
     }
-    new_callback->func = func;
-    new_callback->kwargs = kwargs;
-    Py_INCREF(func);
-    Py_XINCREF(kwargs);
+    callback->func = Py_NewRef(func);
+    callback->kwargs = Py_XNewRef(kwargs);
 
-    modstate->atexit_callbacks[modstate->ncallbacks++] = new_callback;
+    state->atexit_callbacks[state->ncallbacks++] = callback;
 
-    Py_INCREF(func);
-    return func;
+    return Py_NewRef(func);
 }
 
 PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
@@ -187,11 +178,12 @@ PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
 Run all registered exit functions.");
 
 static PyObject *
-atexit_run_exitfuncs(PyObject *self, PyObject *unused)
+atexit_run_exitfuncs(PyObject *module, PyObject *unused)
 {
-    atexit_callfuncs(self);
-    if (PyErr_Occurred())
+    atexit_callfuncs(module);
+    if (PyErr_Occurred()) {
         return NULL;
+    }
     Py_RETURN_NONE;
 }
 
@@ -201,9 +193,9 @@ PyDoc_STRVAR(atexit_clear__doc__,
 Clear the list of previously registered exit functions.");
 
 static PyObject *
-atexit_clear(PyObject *self, PyObject *unused)
+atexit_clear(PyObject *module, PyObject *unused)
 {
-    atexit_cleanup(get_atexit_state(self));
+    atexit_cleanup(get_atexit_state(module));
     Py_RETURN_NONE;
 }
 
@@ -213,25 +205,18 @@ PyDoc_STRVAR(atexit_ncallbacks__doc__,
 Return the number of registered exit functions.");
 
 static PyObject *
-atexit_ncallbacks(PyObject *self, PyObject *unused)
+atexit_ncallbacks(PyObject *module, PyObject *unused)
 {
-    atexitmodule_state *modstate;
-
-    modstate = get_atexit_state(self);
-
-    return PyLong_FromSsize_t(modstate->ncallbacks);
+    struct atexit_state *state = get_atexit_state(module);
+    return PyLong_FromSsize_t(state->ncallbacks);
 }
 
 static int
-atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
+atexit_m_traverse(PyObject *module, visitproc visit, void *arg)
 {
-    int i;
-    atexitmodule_state *modstate;
-
-    modstate = (atexitmodule_state *)PyModule_GetState(self);
-
-    for (i = 0; i < modstate->ncallbacks; i++) {
-        atexit_callback *cb = modstate->atexit_callbacks[i];
+    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
+    for (int i = 0; i < state->ncallbacks; i++) {
+        atexit_callback *cb = state->atexit_callbacks[i];
         if (cb == NULL)
             continue;
         Py_VISIT(cb->func);
@@ -242,21 +227,19 @@ atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
 }
 
 static int
-atexit_m_clear(PyObject *self)
+atexit_m_clear(PyObject *module)
 {
-    atexitmodule_state *modstate;
-    modstate = (atexitmodule_state *)PyModule_GetState(self);
-    atexit_cleanup(modstate);
+    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
+    atexit_cleanup(state);
     return 0;
 }
 
 static void
-atexit_free(PyObject *m)
+atexit_free(PyObject *module)
 {
-    atexitmodule_state *modstate;
-    modstate = (atexitmodule_state *)PyModule_GetState(m);
-    atexit_cleanup(modstate);
-    PyMem_Free(modstate->atexit_callbacks);
+    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
+    atexit_cleanup(state);
+    PyMem_Free(state->atexit_callbacks);
 }
 
 PyDoc_STRVAR(atexit_unregister__doc__,
@@ -268,29 +251,28 @@ atexit.register\n\
     func - function to be unregistered");
 
 static PyObject *
-atexit_unregister(PyObject *self, PyObject *func)
+atexit_unregister(PyObject *module, PyObject *func)
 {
-    atexitmodule_state *modstate;
-    atexit_callback *cb;
-    int i, eq;
-
-    modstate = get_atexit_state(self);
-
-    for (i = 0; i < modstate->ncallbacks; i++)
+    struct atexit_state *state = get_atexit_state(module);
+    for (int i = 0; i < state->ncallbacks; i++)
     {
-        cb = modstate->atexit_callbacks[i];
-        if (cb == NULL)
+        atexit_callback *cb = state->atexit_callbacks[i];
+        if (cb == NULL) {
             continue;
+        }
 
-        eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
-        if (eq < 0)
+        int eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
+        if (eq < 0) {
             return NULL;
-        if (eq)
-            atexit_delete_cb(modstate, i);
+        }
+        if (eq) {
+            atexit_delete_cb(state, i);
+        }
     }
     Py_RETURN_NONE;
 }
 
+
 static PyMethodDef atexit_methods[] = {
     {"register", (PyCFunction)(void(*)(void)) atexit_register, METH_VARARGS|METH_KEYWORDS,
         atexit_register__doc__},
@@ -316,18 +298,19 @@ 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)
+atexit_exec(PyObject *module)
+{
+    struct atexit_state *state = get_atexit_state(module);
+    state->callback_len = 32;
+    state->ncallbacks = 0;
+    state->atexit_callbacks = PyMem_New(atexit_callback*, state->callback_len);
+    if (state->atexit_callbacks == NULL) {
         return -1;
+    }
 
-    _Py_PyAtExit(atexit_callfuncs, m);
+    PyInterpreterState *is = _PyInterpreterState_GET();
+    is->atexit_func = atexit_callfuncs;
+    is->atexit_module = module;
     return 0;
 }
 
@@ -338,14 +321,14 @@ static PyModuleDef_Slot atexit_slots[] = {
 
 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
+    .m_name = "atexit",
+    .m_doc = atexit__doc__,
+    .m_size = sizeof(struct atexit_state),
+    .m_methods = atexit_methods,
+    .m_slots = atexit_slots,
+    .m_traverse = atexit_m_traverse,
+    .m_clear = atexit_m_clear,
+    .m_free = (freefunc)atexit_free
 };
 
 PyMODINIT_FUNC