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
11 /* Forward declaration (for atexit_cleanup) */
12 static PyObject *atexit_clear(PyObject*, PyObject*);
13 /* Forward declaration of module object */
14 static struct PyModuleDef atexitmodule;
15
16 /* ===================================================================== */
17 /* Callback machinery. */
18
19 typedef struct {
20 PyObject *func;
21 PyObject *args;
22 PyObject *kwargs;
23 } atexit_callback;
24
25 typedef struct {
26 atexit_callback **atexit_callbacks;
27 int ncallbacks;
28 int callback_len;
29 } atexitmodule_state;
30
31 static inline atexitmodule_state*
get_atexit_state(PyObject * module)32 get_atexit_state(PyObject *module)
33 {
34 void *state = PyModule_GetState(module);
35 assert(state != NULL);
36 return (atexitmodule_state *)state;
37 }
38
39
40 static void
atexit_delete_cb(atexitmodule_state * modstate,int i)41 atexit_delete_cb(atexitmodule_state *modstate, int i)
42 {
43 atexit_callback *cb;
44
45 cb = modstate->atexit_callbacks[i];
46 modstate->atexit_callbacks[i] = NULL;
47 Py_DECREF(cb->func);
48 Py_DECREF(cb->args);
49 Py_XDECREF(cb->kwargs);
50 PyMem_Free(cb);
51 }
52
53 /* Clear all callbacks without calling them */
54 static void
atexit_cleanup(atexitmodule_state * modstate)55 atexit_cleanup(atexitmodule_state *modstate)
56 {
57 atexit_callback *cb;
58 int i;
59 for (i = 0; i < modstate->ncallbacks; i++) {
60 cb = modstate->atexit_callbacks[i];
61 if (cb == NULL)
62 continue;
63
64 atexit_delete_cb(modstate, i);
65 }
66 modstate->ncallbacks = 0;
67 }
68
69 /* Installed into pylifecycle.c's atexit mechanism */
70
71 static void
atexit_callfuncs(PyObject * module)72 atexit_callfuncs(PyObject *module)
73 {
74 PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;
75 atexit_callback *cb;
76 atexitmodule_state *modstate;
77 int i;
78
79 if (module == NULL)
80 return;
81 modstate = get_atexit_state(module);
82
83 if (modstate->ncallbacks == 0)
84 return;
85
86
87 for (i = modstate->ncallbacks - 1; i >= 0; i--)
88 {
89 cb = modstate->atexit_callbacks[i];
90 if (cb == NULL)
91 continue;
92
93 r = PyObject_Call(cb->func, cb->args, cb->kwargs);
94 Py_XDECREF(r);
95 if (r == NULL) {
96 /* Maintain the last exception, but don't leak if there are
97 multiple exceptions. */
98 if (exc_type) {
99 Py_DECREF(exc_type);
100 Py_XDECREF(exc_value);
101 Py_XDECREF(exc_tb);
102 }
103 PyErr_Fetch(&exc_type, &exc_value, &exc_tb);
104 if (!PyErr_GivenExceptionMatches(exc_type, PyExc_SystemExit)) {
105 PySys_WriteStderr("Error in atexit._run_exitfuncs:\n");
106 PyErr_NormalizeException(&exc_type, &exc_value, &exc_tb);
107 PyErr_Display(exc_type, exc_value, exc_tb);
108 }
109 }
110 }
111
112 atexit_cleanup(modstate);
113
114 if (exc_type)
115 PyErr_Restore(exc_type, exc_value, exc_tb);
116 }
117
118 /* ===================================================================== */
119 /* Module methods. */
120
121 PyDoc_STRVAR(atexit_register__doc__,
122 "register(func, *args, **kwargs) -> func\n\
123 \n\
124 Register a function to be executed upon normal program termination\n\
125 \n\
126 func - function to be called at exit\n\
127 args - optional arguments to pass to func\n\
128 kwargs - optional keyword arguments to pass to func\n\
129 \n\
130 func is returned to facilitate usage as a decorator.");
131
132 static PyObject *
atexit_register(PyObject * self,PyObject * args,PyObject * kwargs)133 atexit_register(PyObject *self, PyObject *args, PyObject *kwargs)
134 {
135 atexitmodule_state *modstate;
136 atexit_callback *new_callback;
137 PyObject *func = NULL;
138
139 modstate = get_atexit_state(self);
140
141 if (modstate->ncallbacks >= modstate->callback_len) {
142 atexit_callback **r;
143 modstate->callback_len += 16;
144 r = (atexit_callback**)PyMem_Realloc(modstate->atexit_callbacks,
145 sizeof(atexit_callback*) * modstate->callback_len);
146 if (r == NULL)
147 return PyErr_NoMemory();
148 modstate->atexit_callbacks = r;
149 }
150
151 if (PyTuple_GET_SIZE(args) == 0) {
152 PyErr_SetString(PyExc_TypeError,
153 "register() takes at least 1 argument (0 given)");
154 return NULL;
155 }
156
157 func = PyTuple_GET_ITEM(args, 0);
158 if (!PyCallable_Check(func)) {
159 PyErr_SetString(PyExc_TypeError,
160 "the first argument must be callable");
161 return NULL;
162 }
163
164 new_callback = PyMem_Malloc(sizeof(atexit_callback));
165 if (new_callback == NULL)
166 return PyErr_NoMemory();
167
168 new_callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
169 if (new_callback->args == NULL) {
170 PyMem_Free(new_callback);
171 return NULL;
172 }
173 new_callback->func = func;
174 new_callback->kwargs = kwargs;
175 Py_INCREF(func);
176 Py_XINCREF(kwargs);
177
178 modstate->atexit_callbacks[modstate->ncallbacks++] = new_callback;
179
180 Py_INCREF(func);
181 return func;
182 }
183
184 PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
185 "_run_exitfuncs() -> None\n\
186 \n\
187 Run all registered exit functions.");
188
189 static PyObject *
atexit_run_exitfuncs(PyObject * self,PyObject * unused)190 atexit_run_exitfuncs(PyObject *self, PyObject *unused)
191 {
192 atexit_callfuncs(self);
193 if (PyErr_Occurred())
194 return NULL;
195 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 *
atexit_clear(PyObject * self,PyObject * unused)204 atexit_clear(PyObject *self, PyObject *unused)
205 {
206 atexit_cleanup(get_atexit_state(self));
207 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 *
atexit_ncallbacks(PyObject * self,PyObject * unused)216 atexit_ncallbacks(PyObject *self, PyObject *unused)
217 {
218 atexitmodule_state *modstate;
219
220 modstate = get_atexit_state(self);
221
222 return PyLong_FromSsize_t(modstate->ncallbacks);
223 }
224
225 static int
atexit_m_traverse(PyObject * self,visitproc visit,void * arg)226 atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
227 {
228 int i;
229 atexitmodule_state *modstate;
230
231 modstate = (atexitmodule_state *)PyModule_GetState(self);
232
233 for (i = 0; i < modstate->ncallbacks; i++) {
234 atexit_callback *cb = modstate->atexit_callbacks[i];
235 if (cb == NULL)
236 continue;
237 Py_VISIT(cb->func);
238 Py_VISIT(cb->args);
239 Py_VISIT(cb->kwargs);
240 }
241 return 0;
242 }
243
244 static int
atexit_m_clear(PyObject * self)245 atexit_m_clear(PyObject *self)
246 {
247 atexitmodule_state *modstate;
248 modstate = (atexitmodule_state *)PyModule_GetState(self);
249 atexit_cleanup(modstate);
250 return 0;
251 }
252
253 static void
atexit_free(PyObject * m)254 atexit_free(PyObject *m)
255 {
256 atexitmodule_state *modstate;
257 modstate = (atexitmodule_state *)PyModule_GetState(m);
258 atexit_cleanup(modstate);
259 PyMem_Free(modstate->atexit_callbacks);
260 }
261
262 PyDoc_STRVAR(atexit_unregister__doc__,
263 "unregister(func) -> None\n\
264 \n\
265 Unregister an exit function which was previously registered using\n\
266 atexit.register\n\
267 \n\
268 func - function to be unregistered");
269
270 static PyObject *
atexit_unregister(PyObject * self,PyObject * func)271 atexit_unregister(PyObject *self, PyObject *func)
272 {
273 atexitmodule_state *modstate;
274 atexit_callback *cb;
275 int i, eq;
276
277 modstate = get_atexit_state(self);
278
279 for (i = 0; i < modstate->ncallbacks; i++)
280 {
281 cb = modstate->atexit_callbacks[i];
282 if (cb == NULL)
283 continue;
284
285 eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
286 if (eq < 0)
287 return NULL;
288 if (eq)
289 atexit_delete_cb(modstate, i);
290 }
291 Py_RETURN_NONE;
292 }
293
294 static PyMethodDef atexit_methods[] = {
295 {"register", (PyCFunction)(void(*)(void)) atexit_register, METH_VARARGS|METH_KEYWORDS,
296 atexit_register__doc__},
297 {"_clear", (PyCFunction) atexit_clear, METH_NOARGS,
298 atexit_clear__doc__},
299 {"unregister", (PyCFunction) atexit_unregister, METH_O,
300 atexit_unregister__doc__},
301 {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS,
302 atexit_run_exitfuncs__doc__},
303 {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS,
304 atexit_ncallbacks__doc__},
305 {NULL, NULL} /* sentinel */
306 };
307
308 /* ===================================================================== */
309 /* Initialization function. */
310
311 PyDoc_STRVAR(atexit__doc__,
312 "allow programmer to define multiple exit functions to be executed\
313 upon normal program termination.\n\
314 \n\
315 Two public functions, register and unregister, are defined.\n\
316 ");
317
318 static int
atexit_exec(PyObject * m)319 atexit_exec(PyObject *m) {
320 atexitmodule_state *modstate;
321
322 modstate = get_atexit_state(m);
323 modstate->callback_len = 32;
324 modstate->ncallbacks = 0;
325 modstate->atexit_callbacks = PyMem_New(atexit_callback*,
326 modstate->callback_len);
327 if (modstate->atexit_callbacks == NULL)
328 return -1;
329
330 _Py_PyAtExit(atexit_callfuncs, m);
331 return 0;
332 }
333
334 static PyModuleDef_Slot atexit_slots[] = {
335 {Py_mod_exec, atexit_exec},
336 {0, NULL}
337 };
338
339 static struct PyModuleDef atexitmodule = {
340 PyModuleDef_HEAD_INIT,
341 "atexit",
342 atexit__doc__,
343 sizeof(atexitmodule_state),
344 atexit_methods,
345 atexit_slots,
346 atexit_m_traverse,
347 atexit_m_clear,
348 (freefunc)atexit_free
349 };
350
351 PyMODINIT_FUNC
PyInit_atexit(void)352 PyInit_atexit(void)
353 {
354 return PyModuleDef_Init(&atexitmodule);
355 }
356