1 files changed, 40 insertions, 32 deletions

diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c
index 35ebf08..afa1cfa 100644
--- a/Modules/atexitmodule.c
+++ b/Modules/atexitmodule.c
@@ -63,15 +63,13 @@ atexit_cleanup(atexitmodule_state *modstate)
 /* Installed into pylifecycle.c's atexit mechanism */
 
 static void
-atexit_callfuncs(void)
+atexit_callfuncs(PyObject *module)
 {
     PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;
     atexit_callback *cb;
-    PyObject *module;
     atexitmodule_state *modstate;
     int i;
 
-    module = PyState_FindModule(&atexitmodule);
     if (module == NULL)
         return;
     modstate = GET_ATEXIT_STATE(module);
@@ -185,7 +183,7 @@ Run all registered exit functions.");
 static PyObject *
 atexit_run_exitfuncs(PyObject *self, PyObject *unused)
 {
-    atexit_callfuncs();
+    atexit_callfuncs(self);
     if (PyErr_Occurred())
         return NULL;
     Py_RETURN_NONE;
@@ -225,13 +223,15 @@ atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
     atexitmodule_state *modstate;
 
     modstate = GET_ATEXIT_STATE(self);
-    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);
+    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;
 }
@@ -241,7 +241,9 @@ atexit_m_clear(PyObject *self)
 {
     atexitmodule_state *modstate;
     modstate = GET_ATEXIT_STATE(self);
-    atexit_cleanup(modstate);
+    if (modstate != NULL) {
+        atexit_cleanup(modstate);
+    }
     return 0;
 }
 
@@ -250,8 +252,10 @@ atexit_free(PyObject *m)
 {
     atexitmodule_state *modstate;
     modstate = GET_ATEXIT_STATE(m);
-    atexit_cleanup(modstate);
-    PyMem_Free(modstate->atexit_callbacks);
+    if (modstate != NULL) {
+        atexit_cleanup(modstate);
+        PyMem_Free(modstate->atexit_callbacks);
+    }
 }
 
 PyDoc_STRVAR(atexit_unregister__doc__,
@@ -310,6 +314,26 @@ upon normal program termination.\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,
@@ -317,7 +341,7 @@ static struct PyModuleDef atexitmodule = {
     atexit__doc__,
     sizeof(atexitmodule_state),
     atexit_methods,
-    NULL,
+    atexit_slots,
     atexit_m_traverse,
     atexit_m_clear,
     (freefunc)atexit_free
@@ -326,21 +350,5 @@ static struct PyModuleDef atexitmodule = {
 PyMODINIT_FUNC
 PyInit_atexit(void)
 {
-    PyObject *m;
-    atexitmodule_state *modstate;
-
-    m = PyModule_Create(&atexitmodule);
-    if (m == NULL)
-        return NULL;
-
-    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 NULL;
-
-    _Py_PyAtExit(atexit_callfuncs);
-    return m;
+    return PyModuleDef_Init(&atexitmodule);
 }