5 files changed, 125 insertions, 158 deletions

diff --git a/Include/cpython/pylifecycle.h b/Include/cpython/pylifecycle.h
index b4e2c8a..13f7a26 100644
--- a/Include/cpython/pylifecycle.h
+++ b/Include/cpython/pylifecycle.h
@@ -35,11 +35,6 @@ PyAPI_FUNC(int) Py_RunMain(void);
 
 PyAPI_FUNC(void) _Py_NO_RETURN Py_ExitStatusException(PyStatus err);
 
-/* Py_PyAtExit is for the atexit module, Py_AtExit is for low-level
- * exit functions.
- */
-PyAPI_FUNC(void) _Py_PyAtExit(void (*func)(PyObject *), PyObject *);
-
 /* Restore signals that the interpreter has called SIG_IGN on to SIG_DFL. */
 PyAPI_FUNC(void) _Py_RestoreSignals(void);
 
diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
index 184878c..2515234 100644
--- a/Include/internal/pycore_interp.h
+++ b/Include/internal/pycore_interp.h
@@ -234,8 +234,8 @@ struct _is {
     PyObject *after_forkers_child;
 #endif
     /* AtExit module */
-    void (*pyexitfunc)(PyObject *);
-    PyObject *pyexitmodule;
+    void (*atexit_func)(PyObject *);
+    PyObject *atexit_module;
 
     uint64_t tstate_next_unique_id;
 
diff --git a/Lib/test/test_atexit.py b/Lib/test/test_atexit.py
index 3105f6c..906f96d 100644
--- a/Lib/test/test_atexit.py
+++ b/Lib/test/test_atexit.py
@@ -1,8 +1,9 @@
-import sys
-import unittest
-import io
 import atexit
+import io
 import os
+import sys
+import textwrap
+import unittest
 from test import support
 from test.support import script_helper
 
@@ -156,7 +157,7 @@ class GeneralTest(unittest.TestCase):
 
     def test_shutdown(self):
         # Actually test the shutdown mechanism in a subprocess
-        code = """if 1:
+        code = textwrap.dedent("""
             import atexit
 
             def f(msg):
@@ -164,7 +165,7 @@ class GeneralTest(unittest.TestCase):
 
             atexit.register(f, "one")
             atexit.register(f, "two")
-            """
+        """)
         res = script_helper.assert_python_ok("-c", code)
         self.assertEqual(res.out.decode().splitlines(), ["two", "one"])
         self.assertFalse(res.err)
@@ -178,13 +179,13 @@ class SubinterpreterTest(unittest.TestCase):
         # take care to free callbacks in its per-subinterpreter module
         # state.
         n = atexit._ncallbacks()
-        code = r"""if 1:
+        code = textwrap.dedent(r"""
             import atexit
             def f():
                 pass
             atexit.register(f)
             del atexit
-            """
+        """)
         ret = support.run_in_subinterp(code)
         self.assertEqual(ret, 0)
         self.assertEqual(atexit._ncallbacks(), n)
@@ -193,13 +194,13 @@ class SubinterpreterTest(unittest.TestCase):
         # Similar to the above, but with a refcycle through the atexit
         # module.
         n = atexit._ncallbacks()
-        code = r"""if 1:
+        code = textwrap.dedent(r"""
             import atexit
             def f():
                 pass
             atexit.register(f)
             atexit.__atexit = atexit
-            """
+        """)
         ret = support.run_in_subinterp(code)
         self.assertEqual(ret, 0)
         self.assertEqual(atexit._ncallbacks(), n)
@@ -210,13 +211,13 @@ class SubinterpreterTest(unittest.TestCase):
         expected = b"The test has passed!"
         r, w = os.pipe()
 
-        code = r"""if 1:
+        code = textwrap.dedent(r"""
             import os
             import atexit
             def callback():
                 os.write({:d}, b"The test has passed!")
             atexit.register(callback)
-        """.format(w)
+        """.format(w))
         ret = support.run_in_subinterp(code)
         os.close(w)
         self.assertEqual(os.read(r, len(expected)), expected)
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
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 6a705b4..54a35a2 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -2634,26 +2634,14 @@ Py_ExitStatusException(PyStatus status)
 
 /* Clean up and exit */
 
-/* For the atexit module. */
-void _Py_PyAtExit(void (*func)(PyObject *), PyObject *module)
-{
-    PyInterpreterState *is = _PyInterpreterState_GET();
-
-    /* Guard against API misuse (see bpo-17852) */
-    assert(is->pyexitfunc == NULL || is->pyexitfunc == func);
-
-    is->pyexitfunc = func;
-    is->pyexitmodule = module;
-}
-
 static void
 call_py_exitfuncs(PyThreadState *tstate)
 {
     PyInterpreterState *interp = tstate->interp;
-    if (interp->pyexitfunc == NULL)
+    if (interp->atexit_func == NULL)
         return;
 
-    (*interp->pyexitfunc)(interp->pyexitmodule);
+    interp->atexit_func(interp->atexit_module);
     _PyErr_Clear(tstate);
 }