bpo-31356: Add context manager to temporarily disable GC (GH-4224) · python/cpython@72a0d21
@@ -1067,6 +1067,10 @@ static PyObject *
10671067gc_enable_impl(PyObject *module)
10681068/*[clinic end generated code: output=45a427e9dce9155c input=81ac4940ca579707]*/
10691069{
1070+if(_PyRuntime.gc.disabled_threads){
1071+PyErr_WarnEx(PyExc_RuntimeWarning, "Garbage collector enabled while another "
1072+"thread is inside gc.ensure_enabled",1);
1073+ }
10701074_PyRuntime.gc.enabled = 1;
10711075Py_RETURN_NONE;
10721076}
@@ -1508,6 +1512,102 @@ static PyMethodDef GcMethods[] = {
15081512 {NULL, NULL} /* Sentinel */
15091513};
151015141515+typedef struct {
1516+PyObject_HEAD
1517+int previous_gc_state;
1518+} ensure_disabled_object;
1519+1520+1521+static void
1522+ensure_disabled_object_dealloc(ensure_disabled_object *m_obj)
1523+{
1524+Py_TYPE(m_obj)->tp_free((PyObject*)m_obj);
1525+}
1526+1527+static PyObject *
1528+ensure_disabled__enter__method(ensure_disabled_object *self, PyObject *args)
1529+{
1530+PyGILState_STATE gstate = PyGILState_Ensure();
1531+++_PyRuntime.gc.disabled_threads;
1532+self->previous_gc_state = _PyRuntime.gc.enabled;
1533+gc_disable_impl(NULL);
1534+PyGILState_Release(gstate);
1535+Py_RETURN_NONE;
1536+}
1537+1538+static PyObject *
1539+ensure_disabled__exit__method(ensure_disabled_object *self, PyObject *args)
1540+{
1541+PyGILState_STATE gstate = PyGILState_Ensure();
1542+--_PyRuntime.gc.disabled_threads;
1543+if(self->previous_gc_state){
1544+gc_enable_impl(NULL);
1545+ }else{
1546+gc_disable_impl(NULL);
1547+ }
1548+PyGILState_Release(gstate);
1549+Py_RETURN_NONE;
1550+}
1551+1552+1553+1554+static struct PyMethodDef ensure_disabled_object_methods[] = {
1555+ {"__enter__", (PyCFunction) ensure_disabled__enter__method, METH_NOARGS},
1556+ {"__exit__", (PyCFunction) ensure_disabled__exit__method, METH_VARARGS},
1557+ {NULL, NULL} /* sentinel */
1558+};
1559+1560+static PyObject *
1561+new_disabled_obj(PyTypeObject *type, PyObject *args, PyObject *kwdict){
1562+ensure_disabled_object *self;
1563+self = (ensure_disabled_object *)type->tp_alloc(type, 0);
1564+return (PyObject *) self;
1565+};
1566+1567+static PyTypeObject gc_ensure_disabled_type = {
1568+PyVarObject_HEAD_INIT(NULL, 0)
1569+"gc.ensure_disabled", /* tp_name */
1570+sizeof(ensure_disabled_object), /* tp_size */
1571+0, /* tp_itemsize */
1572+/* methods */
1573+ (destructor) ensure_disabled_object_dealloc,/* tp_dealloc */
1574+0, /* tp_print */
1575+0, /* tp_getattr */
1576+0, /* tp_setattr */
1577+0, /* tp_reserved */
1578+0, /* tp_repr */
1579+0, /* tp_as_number */
1580+0, /*tp_as_sequence*/
1581+0, /*tp_as_mapping*/
1582+0, /*tp_hash*/
1583+0, /*tp_call*/
1584+0, /*tp_str*/
1585+PyObject_GenericGetAttr, /*tp_getattro*/
1586+0, /*tp_setattro*/
1587+0, /*tp_as_buffer*/
1588+Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
1589+0, /*tp_doc*/
1590+0, /* tp_traverse */
1591+0, /* tp_clear */
1592+0, /* tp_richcompare */
1593+0, /* tp_weaklistoffset */
1594+0, /* tp_iter */
1595+0, /* tp_iternext */
1596+ensure_disabled_object_methods, /* tp_methods */
1597+0, /* tp_members */
1598+0, /* tp_getset */
1599+0, /* tp_base */
1600+0, /* tp_dict */
1601+0, /* tp_descr_get */
1602+0, /* tp_descr_set */
1603+0, /* tp_dictoffset */
1604+0, /* tp_init */
1605+PyType_GenericAlloc, /* tp_alloc */
1606+new_disabled_obj, /* tp_new */
1607+PyObject_Del, /* tp_free */
1608+};
1609+1610+15111611static struct PyModuleDef gcmodule = {
15121612PyModuleDef_HEAD_INIT,
15131613"gc", /* m_name */
@@ -1548,6 +1648,12 @@ PyInit_gc(void)
15481648if (PyModule_AddObject(m, "callbacks", _PyRuntime.gc.callbacks) < 0)
15491649return NULL;
155016501651+if (PyType_Ready(&gc_ensure_disabled_type) < 0)
1652+return NULL;
1653+if (PyModule_AddObject(m, "ensure_disabled", (PyObject*) &gc_ensure_disabled_type) < 0)
1654+return NULL;
1655+1656+15511657#define ADD_INT(NAME) if (PyModule_AddIntConstant(m, #NAME, NAME) < 0) return NULL
15521658ADD_INT(DEBUG_STATS);
15531659ADD_INT(DEBUG_COLLECTABLE);