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 #include "pycore_initconfig.h" // _PyStatus_NO_MEMORY
11 #include "pycore_interp.h" // PyInterpreterState.atexit
12 #include "pycore_pystate.h" // _PyInterpreterState_GET
13
14 /* ===================================================================== */
15 /* Callback machinery. */
16
17 static inline struct atexit_state*
18 get_atexit_state(void)
19 {
20 PyInterpreterState *interp = _PyInterpreterState_GET();
21 return &interp->atexit;
22 }
23
24
25 static void
26 atexit_delete_cb(struct atexit_state *state, int i)
27 {
28 atexit_callback *cb = state->callbacks[i];
29 state->callbacks[i] = NULL;
30
31 Py_DECREF(cb->func);
32 Py_DECREF(cb->args);
33 Py_XDECREF(cb->kwargs);
34 PyMem_Free(cb);
35 }
36
37
38 /* Clear all callbacks without calling them */
39 static void
40 atexit_cleanup(struct atexit_state *state)
41 {
42 atexit_callback *cb;
43 for (int i = 0; i < state->ncallbacks; i++) {
44 cb = state->callbacks[i];
45 if (cb == NULL)
46 continue;
47
48 atexit_delete_cb(state, i);
49 }
50 state->ncallbacks = 0;
51 }
52
53
54 PyStatus
55 _PyAtExit_Init(PyInterpreterState *interp)
56 {
57 struct atexit_state *state = &interp->atexit;
58 // _PyAtExit_Init() must only be called once
59 assert(state->callbacks == NULL);
60
61 state->callback_len = 32;
62 state->ncallbacks = 0;
63 state->callbacks = PyMem_New(atexit_callback*, state->callback_len);
64 if (state->callbacks == NULL) {
65 return _PyStatus_NO_MEMORY();
66 }
67 return _PyStatus_OK();
68 }
69
70
71 void
72 _PyAtExit_Fini(PyInterpreterState *interp)
73 {
74 struct atexit_state *state = &interp->atexit;
75 atexit_cleanup(state);
76 PyMem_Free(state->callbacks);
77 state->callbacks = NULL;
78 }
79
80
81 static void
82 atexit_callfuncs(struct atexit_state *state)
83 {
84 assert(!PyErr_Occurred());
85
86 if (state->ncallbacks == 0) {
87 return;
88 }
89
90 for (int i = state->ncallbacks - 1; i >= 0; i--) {
91 atexit_callback *cb = state->callbacks[i];
92 if (cb == NULL) {
93 continue;
94 }
95
96 // bpo-46025: Increment the refcount of cb->func as the call itself may unregister it
97 PyObject* the_func = Py_NewRef(cb->func);
98 PyObject *res = PyObject_Call(cb->func, cb->args, cb->kwargs);
99 if (res == NULL) {
100 _PyErr_WriteUnraisableMsg("in atexit callback", the_func);
101 }
102 else {
103 Py_DECREF(res);
104 }
105 Py_DECREF(the_func);
106 }
107
108 atexit_cleanup(state);
109
110 assert(!PyErr_Occurred());
111 }
112
113
114 void
115 _PyAtExit_Call(PyInterpreterState *interp)
116 {
117 struct atexit_state *state = &interp->atexit;
118 atexit_callfuncs(state);
119 }
120
121
122 /* ===================================================================== */
123 /* Module methods. */
124
125
126 PyDoc_STRVAR(atexit_register__doc__,
127 "register(func, *args, **kwargs) -> func\n\
128 \n\
129 Register a function to be executed upon normal program termination\n\
130 \n\
131 func - function to be called at exit\n\
132 args - optional arguments to pass to func\n\
133 kwargs - optional keyword arguments to pass to func\n\
134 \n\
135 func is returned to facilitate usage as a decorator.");
136
137 static PyObject *
138 atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)
139 {
140 if (PyTuple_GET_SIZE(args) == 0) {
141 PyErr_SetString(PyExc_TypeError,
142 "register() takes at least 1 argument (0 given)");
143 return NULL;
144 }
145
146 PyObject *func = PyTuple_GET_ITEM(args, 0);
147 if (!PyCallable_Check(func)) {
148 PyErr_SetString(PyExc_TypeError,
149 "the first argument must be callable");
150 return NULL;
151 }
152
153 struct atexit_state *state = get_atexit_state();
154 if (state->ncallbacks >= state->callback_len) {
155 atexit_callback **r;
156 state->callback_len += 16;
157 size_t size = sizeof(atexit_callback*) * (size_t)state->callback_len;
158 r = (atexit_callback**)PyMem_Realloc(state->callbacks, size);
159 if (r == NULL) {
160 return PyErr_NoMemory();
161 }
162 state->callbacks = r;
163 }
164
165 atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
166 if (callback == NULL) {
167 return PyErr_NoMemory();
168 }
169
170 callback->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
171 if (callback->args == NULL) {
172 PyMem_Free(callback);
173 return NULL;
174 }
175 callback->func = Py_NewRef(func);
176 callback->kwargs = Py_XNewRef(kwargs);
177
178 state->callbacks[state->ncallbacks++] = callback;
179
180 return Py_NewRef(func);
181 }
182
183 PyDoc_STRVAR(atexit_run_exitfuncs__doc__,
184 "_run_exitfuncs() -> None\n\
185 \n\
186 Run all registered exit functions.\n\
187 \n\
188 If a callback raises an exception, it is logged with sys.unraisablehook.");
189
190 static PyObject *
191 atexit_run_exitfuncs(PyObject *module, PyObject *unused)
192 {
193 struct atexit_state *state = get_atexit_state();
194 atexit_callfuncs(state);
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 *
204 atexit_clear(PyObject *module, PyObject *unused)
205 {
206 atexit_cleanup(get_atexit_state());
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 *
216 atexit_ncallbacks(PyObject *module, PyObject *unused)
217 {
218 struct atexit_state *state = get_atexit_state();
219 return PyLong_FromSsize_t(state->ncallbacks);
220 }
221
222 PyDoc_STRVAR(atexit_unregister__doc__,
223 "unregister(func) -> None\n\
224 \n\
225 Unregister an exit function which was previously registered using\n\
226 atexit.register\n\
227 \n\
228 func - function to be unregistered");
229
230 static PyObject *
231 atexit_unregister(PyObject *module, PyObject *func)
232 {
233 struct atexit_state *state = get_atexit_state();
234 for (int i = 0; i < state->ncallbacks; i++)
235 {
236 atexit_callback *cb = state->callbacks[i];
237 if (cb == NULL) {
238 continue;
239 }
240
241 int eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
242 if (eq < 0) {
243 return NULL;
244 }
245 if (eq) {
246 atexit_delete_cb(state, i);
247 }
248 }
249 Py_RETURN_NONE;
250 }
251
252
253 static PyMethodDef atexit_methods[] = {
254 {"register", _PyCFunction_CAST(atexit_register), METH_VARARGS|METH_KEYWORDS,
255 atexit_register__doc__},
256 {"_clear", (PyCFunction) atexit_clear, METH_NOARGS,
257 atexit_clear__doc__},
258 {"unregister", (PyCFunction) atexit_unregister, METH_O,
259 atexit_unregister__doc__},
260 {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS,
261 atexit_run_exitfuncs__doc__},
262 {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS,
263 atexit_ncallbacks__doc__},
264 {NULL, NULL} /* sentinel */
265 };
266
267
268 /* ===================================================================== */
269 /* Initialization function. */
270
271 PyDoc_STRVAR(atexit__doc__,
272 "allow programmer to define multiple exit functions to be executed\n\
273 upon normal program termination.\n\
274 \n\
275 Two public functions, register and unregister, are defined.\n\
276 ");
277
278 static struct PyModuleDef atexitmodule = {
279 PyModuleDef_HEAD_INIT,
280 .m_name = "atexit",
281 .m_doc = atexit__doc__,
282 .m_size = 0,
283 .m_methods = atexit_methods,
284 };
285
286 PyMODINIT_FUNC
287 PyInit_atexit(void)
288 {
289 return PyModuleDef_Init(&atexitmodule);
290 }