1 #include "parts.h"
2
3 static PyObject*
test_gc_control(PyObject * self,PyObject * Py_UNUSED (ignored))4 test_gc_control(PyObject *self, PyObject *Py_UNUSED(ignored))
5 {
6 int orig_enabled = PyGC_IsEnabled();
7 const char* msg = "ok";
8 int old_state;
9
10 old_state = PyGC_Enable();
11 msg = "Enable(1)";
12 if (old_state != orig_enabled) {
13 goto failed;
14 }
15 msg = "IsEnabled(1)";
16 if (!PyGC_IsEnabled()) {
17 goto failed;
18 }
19
20 old_state = PyGC_Disable();
21 msg = "disable(2)";
22 if (!old_state) {
23 goto failed;
24 }
25 msg = "IsEnabled(2)";
26 if (PyGC_IsEnabled()) {
27 goto failed;
28 }
29
30 old_state = PyGC_Enable();
31 msg = "enable(3)";
32 if (old_state) {
33 goto failed;
34 }
35 msg = "IsEnabled(3)";
36 if (!PyGC_IsEnabled()) {
37 goto failed;
38 }
39
40 if (!orig_enabled) {
41 old_state = PyGC_Disable();
42 msg = "disable(4)";
43 if (old_state) {
44 goto failed;
45 }
46 msg = "IsEnabled(4)";
47 if (PyGC_IsEnabled()) {
48 goto failed;
49 }
50 }
51
52 Py_RETURN_NONE;
53
54 failed:
55 /* Try to clean up if we can. */
56 if (orig_enabled) {
57 PyGC_Enable();
58 } else {
59 PyGC_Disable();
60 }
61 PyErr_Format(PyExc_ValueError, "GC control failed in %s", msg);
62 return NULL;
63 }
64
65 static PyObject *
without_gc(PyObject * Py_UNUSED (self),PyObject * obj)66 without_gc(PyObject *Py_UNUSED(self), PyObject *obj)
67 {
68 PyTypeObject *tp = (PyTypeObject*)obj;
69 if (!PyType_Check(obj) || !PyType_HasFeature(tp, Py_TPFLAGS_HEAPTYPE)) {
70 return PyErr_Format(PyExc_TypeError, "heap type expected, got %R", obj);
71 }
72 if (PyType_IS_GC(tp)) {
73 // Don't try this at home, kids:
74 tp->tp_flags -= Py_TPFLAGS_HAVE_GC;
75 tp->tp_free = PyObject_Del;
76 tp->tp_traverse = NULL;
77 tp->tp_clear = NULL;
78 }
79 assert(!PyType_IS_GC(tp));
80 return Py_NewRef(obj);
81 }
82
83 static void
slot_tp_del(PyObject * self)84 slot_tp_del(PyObject *self)
85 {
86 PyObject *del, *res;
87
88 /* Temporarily resurrect the object. */
89 assert(Py_REFCNT(self) == 0);
90 Py_SET_REFCNT(self, 1);
91
92 /* Save the current exception, if any. */
93 PyObject *exc = PyErr_GetRaisedException();
94
95 PyObject *tp_del = PyUnicode_InternFromString("__tp_del__");
96 if (tp_del == NULL) {
97 PyErr_WriteUnraisable(NULL);
98 PyErr_SetRaisedException(exc);
99 return;
100 }
101 /* Execute __del__ method, if any. */
102 del = _PyType_LookupRef(Py_TYPE(self), tp_del);
103 Py_DECREF(tp_del);
104 if (del != NULL) {
105 res = PyObject_CallOneArg(del, self);
106 Py_DECREF(del);
107 if (res == NULL)
108 PyErr_WriteUnraisable(del);
109 else
110 Py_DECREF(res);
111 }
112
113 /* Restore the saved exception. */
114 PyErr_SetRaisedException(exc);
115
116 /* Undo the temporary resurrection; can't use DECREF here, it would
117 * cause a recursive call.
118 */
119 assert(Py_REFCNT(self) > 0);
120 Py_SET_REFCNT(self, Py_REFCNT(self) - 1);
121 if (Py_REFCNT(self) == 0) {
122 /* this is the normal path out */
123 return;
124 }
125
126 /* __del__ resurrected it! Make it look like the original Py_DECREF
127 * never happened.
128 */
129 {
130 _Py_ResurrectReference(self);
131 }
132 assert(!PyType_IS_GC(Py_TYPE(self)) || PyObject_GC_IsTracked(self));
133 }
134
135 static PyObject *
with_tp_del(PyObject * self,PyObject * args)136 with_tp_del(PyObject *self, PyObject *args)
137 {
138 PyObject *obj;
139 PyTypeObject *tp;
140
141 if (!PyArg_ParseTuple(args, "O:with_tp_del", &obj))
142 return NULL;
143 tp = (PyTypeObject *) obj;
144 if (!PyType_Check(obj) || !PyType_HasFeature(tp, Py_TPFLAGS_HEAPTYPE)) {
145 PyErr_Format(PyExc_TypeError,
146 "heap type expected, got %R", obj);
147 return NULL;
148 }
149 tp->tp_del = slot_tp_del;
150 return Py_NewRef(obj);
151 }
152
153
154 struct gc_visit_state_basic {
155 PyObject *target;
156 int found;
157 };
158
159 static int
gc_visit_callback_basic(PyObject * obj,void * arg)160 gc_visit_callback_basic(PyObject *obj, void *arg)
161 {
162 struct gc_visit_state_basic *state = (struct gc_visit_state_basic *)arg;
163 if (obj == state->target) {
164 state->found = 1;
165 return 0;
166 }
167 return 1;
168 }
169
170 static PyObject *
test_gc_visit_objects_basic(PyObject * Py_UNUSED (self),PyObject * Py_UNUSED (ignored))171 test_gc_visit_objects_basic(PyObject *Py_UNUSED(self),
172 PyObject *Py_UNUSED(ignored))
173 {
174 PyObject *obj;
175 struct gc_visit_state_basic state;
176
177 obj = PyList_New(0);
178 if (obj == NULL) {
179 return NULL;
180 }
181 state.target = obj;
182 state.found = 0;
183
184 PyUnstable_GC_VisitObjects(gc_visit_callback_basic, &state);
185 Py_DECREF(obj);
186 if (!state.found) {
187 PyErr_SetString(
188 PyExc_AssertionError,
189 "test_gc_visit_objects_basic: Didn't find live list");
190 return NULL;
191 }
192 Py_RETURN_NONE;
193 }
194
195 static int
gc_visit_callback_exit_early(PyObject * obj,void * arg)196 gc_visit_callback_exit_early(PyObject *obj, void *arg)
197 {
198 int *visited_i = (int *)arg;
199 (*visited_i)++;
200 if (*visited_i == 2) {
201 return 0;
202 }
203 return 1;
204 }
205
206 static PyObject *
test_gc_visit_objects_exit_early(PyObject * Py_UNUSED (self),PyObject * Py_UNUSED (ignored))207 test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
208 PyObject *Py_UNUSED(ignored))
209 {
210 int visited_i = 0;
211 PyUnstable_GC_VisitObjects(gc_visit_callback_exit_early, &visited_i);
212 if (visited_i != 2) {
213 PyErr_SetString(
214 PyExc_AssertionError,
215 "test_gc_visit_objects_exit_early: did not exit when expected");
216 }
217 Py_RETURN_NONE;
218 }
219
220 typedef struct {
221 PyObject_HEAD
222 } ObjExtraData;
223
224 static PyObject *
obj_extra_data_new(PyTypeObject * type,PyObject * args,PyObject * kwds)225 obj_extra_data_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
226 {
227 size_t extra_size = sizeof(PyObject *);
228 PyObject *obj = PyUnstable_Object_GC_NewWithExtraData(type, extra_size);
229 if (obj == NULL) {
230 return PyErr_NoMemory();
231 }
232 PyObject_GC_Track(obj);
233 return obj;
234 }
235
236 static PyObject **
obj_extra_data_get_extra_storage(PyObject * self)237 obj_extra_data_get_extra_storage(PyObject *self)
238 {
239 return (PyObject **)((char *)self + Py_TYPE(self)->tp_basicsize);
240 }
241
242 static PyObject *
obj_extra_data_get(PyObject * self,void * Py_UNUSED (ignored))243 obj_extra_data_get(PyObject *self, void *Py_UNUSED(ignored))
244 {
245 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
246 PyObject *value = *extra_storage;
247 if (!value) {
248 Py_RETURN_NONE;
249 }
250 return Py_NewRef(value);
251 }
252
253 static int
obj_extra_data_set(PyObject * self,PyObject * newval,void * Py_UNUSED (ignored))254 obj_extra_data_set(PyObject *self, PyObject *newval, void *Py_UNUSED(ignored))
255 {
256 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
257 Py_CLEAR(*extra_storage);
258 if (newval) {
259 *extra_storage = Py_NewRef(newval);
260 }
261 return 0;
262 }
263
264 static PyGetSetDef obj_extra_data_getset[] = {
265 {"extra", (getter)obj_extra_data_get, (setter)obj_extra_data_set, NULL},
266 {NULL}
267 };
268
269 static int
obj_extra_data_traverse(PyObject * self,visitproc visit,void * arg)270 obj_extra_data_traverse(PyObject *self, visitproc visit, void *arg)
271 {
272 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
273 PyObject *value = *extra_storage;
274 Py_VISIT(value);
275 return 0;
276 }
277
278 static int
obj_extra_data_clear(PyObject * self)279 obj_extra_data_clear(PyObject *self)
280 {
281 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
282 Py_CLEAR(*extra_storage);
283 return 0;
284 }
285
286 static void
obj_extra_data_dealloc(PyObject * self)287 obj_extra_data_dealloc(PyObject *self)
288 {
289 PyTypeObject *tp = Py_TYPE(self);
290 PyObject_GC_UnTrack(self);
291 obj_extra_data_clear(self);
292 tp->tp_free(self);
293 Py_DECREF(tp);
294 }
295
296 static PyType_Slot ObjExtraData_Slots[] = {
297 {Py_tp_getset, obj_extra_data_getset},
298 {Py_tp_dealloc, obj_extra_data_dealloc},
299 {Py_tp_traverse, obj_extra_data_traverse},
300 {Py_tp_clear, obj_extra_data_clear},
301 {Py_tp_new, obj_extra_data_new},
302 {Py_tp_free, PyObject_GC_Del},
303 {0, NULL},
304 };
305
306 static PyType_Spec ObjExtraData_TypeSpec = {
307 .name = "_testcapi.ObjExtraData",
308 .basicsize = sizeof(ObjExtraData),
309 .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
310 .slots = ObjExtraData_Slots,
311 };
312
313 static PyMethodDef test_methods[] = {
314 {"test_gc_control", test_gc_control, METH_NOARGS},
315 {"test_gc_visit_objects_basic", test_gc_visit_objects_basic, METH_NOARGS, NULL},
316 {"test_gc_visit_objects_exit_early", test_gc_visit_objects_exit_early, METH_NOARGS, NULL},
317 {"without_gc", without_gc, METH_O, NULL},
318 {"with_tp_del", with_tp_del, METH_VARARGS, NULL},
319 {NULL}
320 };
321
_PyTestCapi_Init_GC(PyObject * mod)322 int _PyTestCapi_Init_GC(PyObject *mod)
323 {
324 if (PyModule_AddFunctions(mod, test_methods) < 0) {
325 return -1;
326 }
327 if (PyModule_AddFunctions(mod, test_methods) < 0) {
328 return -1;
329 }
330
331 PyObject *ObjExtraData_Type = PyType_FromModuleAndSpec(
332 mod, &ObjExtraData_TypeSpec, NULL);
333 if (ObjExtraData_Type == 0) {
334 return -1;
335 }
336 int ret = PyModule_AddType(mod, (PyTypeObject*)ObjExtraData_Type);
337 Py_DECREF(ObjExtraData_Type);
338 if (ret < 0) {
339 return ret;
340 }
341
342 return 0;
343 }
344