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