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