1 #include "parts.h"
2
3 static PyObject*
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 *
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
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_Lookup(Py_TYPE(self), tp_del);
103 Py_DECREF(tp_del);
104 if (del != NULL) {
105 res = PyObject_CallOneArg(del, self);
106 if (res == NULL)
107 PyErr_WriteUnraisable(del);
108 else
109 Py_DECREF(res);
110 }
111
112 /* Restore the saved exception. */
113 PyErr_SetRaisedException(exc);
114
115 /* Undo the temporary resurrection; can't use DECREF here, it would
116 * cause a recursive call.
117 */
118 assert(Py_REFCNT(self) > 0);
119 Py_SET_REFCNT(self, Py_REFCNT(self) - 1);
120 if (Py_REFCNT(self) == 0) {
121 /* this is the normal path out */
122 return;
123 }
124
125 /* __del__ resurrected it! Make it look like the original Py_DECREF
126 * never happened.
127 */
128 {
129 Py_ssize_t refcnt = Py_REFCNT(self);
130 _Py_NewReferenceNoTotal(self);
131 Py_SET_REFCNT(self, refcnt);
132 }
133 assert(!PyType_IS_GC(Py_TYPE(self)) || PyObject_GC_IsTracked(self));
134 }
135
136 static PyObject *
137 with_tp_del(PyObject *self, PyObject *args)
138 {
139 PyObject *obj;
140 PyTypeObject *tp;
141
142 if (!PyArg_ParseTuple(args, "O:with_tp_del", &obj))
143 return NULL;
144 tp = (PyTypeObject *) obj;
145 if (!PyType_Check(obj) || !PyType_HasFeature(tp, Py_TPFLAGS_HEAPTYPE)) {
146 PyErr_Format(PyExc_TypeError,
147 "heap type expected, got %R", obj);
148 return NULL;
149 }
150 tp->tp_del = slot_tp_del;
151 return Py_NewRef(obj);
152 }
153
154
155 struct gc_visit_state_basic {
156 PyObject *target;
157 int found;
158 };
159
160 static int
161 gc_visit_callback_basic(PyObject *obj, void *arg)
162 {
163 struct gc_visit_state_basic *state = (struct gc_visit_state_basic *)arg;
164 if (obj == state->target) {
165 state->found = 1;
166 return 0;
167 }
168 return 1;
169 }
170
171 static PyObject *
172 test_gc_visit_objects_basic(PyObject *Py_UNUSED(self),
173 PyObject *Py_UNUSED(ignored))
174 {
175 PyObject *obj;
176 struct gc_visit_state_basic state;
177
178 obj = PyList_New(0);
179 if (obj == NULL) {
180 return NULL;
181 }
182 state.target = obj;
183 state.found = 0;
184
185 PyUnstable_GC_VisitObjects(gc_visit_callback_basic, &state);
186 Py_DECREF(obj);
187 if (!state.found) {
188 PyErr_SetString(
189 PyExc_AssertionError,
190 "test_gc_visit_objects_basic: Didn't find live list");
191 return NULL;
192 }
193 Py_RETURN_NONE;
194 }
195
196 static int
197 gc_visit_callback_exit_early(PyObject *obj, void *arg)
198 {
199 int *visited_i = (int *)arg;
200 (*visited_i)++;
201 if (*visited_i == 2) {
202 return 0;
203 }
204 return 1;
205 }
206
207 static PyObject *
208 test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
209 PyObject *Py_UNUSED(ignored))
210 {
211 int visited_i = 0;
212 PyUnstable_GC_VisitObjects(gc_visit_callback_exit_early, &visited_i);
213 if (visited_i != 2) {
214 PyErr_SetString(
215 PyExc_AssertionError,
216 "test_gc_visit_objects_exit_early: did not exit when expected");
217 }
218 Py_RETURN_NONE;
219 }
220
221 typedef struct {
222 PyObject_HEAD
223 } ObjExtraData;
224
225 static PyObject *
226 obj_extra_data_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
227 {
228 size_t extra_size = sizeof(PyObject *);
229 PyObject *obj = PyUnstable_Object_GC_NewWithExtraData(type, extra_size);
230 if (obj == NULL) {
231 return PyErr_NoMemory();
232 }
233 PyObject_GC_Track(obj);
234 return obj;
235 }
236
237 static PyObject **
238 obj_extra_data_get_extra_storage(PyObject *self)
239 {
240 return (PyObject **)((char *)self + Py_TYPE(self)->tp_basicsize);
241 }
242
243 static PyObject *
244 obj_extra_data_get(PyObject *self, void *Py_UNUSED(ignored))
245 {
246 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
247 PyObject *value = *extra_storage;
248 if (!value) {
249 Py_RETURN_NONE;
250 }
251 return Py_NewRef(value);
252 }
253
254 static int
255 obj_extra_data_set(PyObject *self, PyObject *newval, void *Py_UNUSED(ignored))
256 {
257 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
258 Py_CLEAR(*extra_storage);
259 if (newval) {
260 *extra_storage = Py_NewRef(newval);
261 }
262 return 0;
263 }
264
265 static PyGetSetDef obj_extra_data_getset[] = {
266 {"extra", (getter)obj_extra_data_get, (setter)obj_extra_data_set, NULL},
267 {NULL}
268 };
269
270 static int
271 obj_extra_data_traverse(PyObject *self, visitproc visit, void *arg)
272 {
273 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
274 PyObject *value = *extra_storage;
275 Py_VISIT(value);
276 return 0;
277 }
278
279 static int
280 obj_extra_data_clear(PyObject *self)
281 {
282 PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
283 Py_CLEAR(*extra_storage);
284 return 0;
285 }
286
287 static void
288 obj_extra_data_dealloc(PyObject *self)
289 {
290 PyTypeObject *tp = Py_TYPE(self);
291 PyObject_GC_UnTrack(self);
292 obj_extra_data_clear(self);
293 tp->tp_free(self);
294 Py_DECREF(tp);
295 }
296
297 static PyType_Slot ObjExtraData_Slots[] = {
298 {Py_tp_getset, obj_extra_data_getset},
299 {Py_tp_dealloc, obj_extra_data_dealloc},
300 {Py_tp_traverse, obj_extra_data_traverse},
301 {Py_tp_clear, obj_extra_data_clear},
302 {Py_tp_new, obj_extra_data_new},
303 {Py_tp_free, PyObject_GC_Del},
304 {0, NULL},
305 };
306
307 static PyType_Spec ObjExtraData_TypeSpec = {
308 .name = "_testcapi.ObjExtraData",
309 .basicsize = sizeof(ObjExtraData),
310 .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
311 .slots = ObjExtraData_Slots,
312 };
313
314 static PyMethodDef test_methods[] = {
315 {"test_gc_control", test_gc_control, METH_NOARGS},
316 {"test_gc_visit_objects_basic", test_gc_visit_objects_basic, METH_NOARGS, NULL},
317 {"test_gc_visit_objects_exit_early", test_gc_visit_objects_exit_early, METH_NOARGS, NULL},
318 {"without_gc", without_gc, METH_O, NULL},
319 {"with_tp_del", with_tp_del, METH_VARARGS, NULL},
320 {NULL}
321 };
322
323 int _PyTestCapi_Init_GC(PyObject *mod)
324 {
325 if (PyModule_AddFunctions(mod, test_methods) < 0) {
326 return -1;
327 }
328 if (PyModule_AddFunctions(mod, test_methods) < 0) {
329 return -1;
330 }
331
332 PyObject *ObjExtraData_Type = PyType_FromModuleAndSpec(
333 mod, &ObjExtraData_TypeSpec, NULL);
334 if (ObjExtraData_Type == 0) {
335 return -1;
336 }
337 int ret = PyModule_AddType(mod, (PyTypeObject*)ObjExtraData_Type);
338 Py_DECREF(ObjExtraData_Type);
339 if (ret < 0) {
340 return ret;
341 }
342
343 return 0;
344 }