diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c index dd93acdbe9267..99a0170eb7627 100644 --- a/mypyc/lib-rt/dict_ops.c +++ b/mypyc/lib-rt/dict_ops.c @@ -346,16 +346,24 @@ tuple_T3CIO CPyDict_NextKey(PyObject *dict_or_iter, CPyTagged offset) { PyObject *dummy; if (PyDict_CheckExact(dict_or_iter)) { - ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &ret.f2, &dummy); + PyObject *key; + // PyDict_Next() returns a borrowed reference. On free-threaded builds, + // hold the dict lock until we have converted it to a strong reference. + Py_BEGIN_CRITICAL_SECTION(dict_or_iter); + ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &key, &dummy); + if (ret.f0) { + ret.f2 = Py_NewRef(key); + } + Py_END_CRITICAL_SECTION(); + if (ret.f0) { ret.f1 = CPyTagged_FromSsize_t(py_offset); } else { // Set key to None, so mypyc can manage refcounts. ret.f1 = 0; ret.f2 = Py_None; + Py_INCREF(ret.f2); } - // PyDict_Next() returns borrowed references. - Py_INCREF(ret.f2); } else { // offset is dummy in this case, just use the old value. ret.f1 = offset; @@ -370,16 +378,24 @@ tuple_T3CIO CPyDict_NextValue(PyObject *dict_or_iter, CPyTagged offset) { PyObject *dummy; if (PyDict_CheckExact(dict_or_iter)) { - ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &dummy, &ret.f2); + PyObject *value; + // PyDict_Next() returns a borrowed reference. On free-threaded builds, + // hold the dict lock until we have converted it to a strong reference. + Py_BEGIN_CRITICAL_SECTION(dict_or_iter); + ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &dummy, &value); + if (ret.f0) { + ret.f2 = Py_NewRef(value); + } + Py_END_CRITICAL_SECTION(); + if (ret.f0) { ret.f1 = CPyTagged_FromSsize_t(py_offset); } else { // Set value to None, so mypyc can manage refcounts. ret.f1 = 0; ret.f2 = Py_None; + Py_INCREF(ret.f2); } - // PyDict_Next() returns borrowed references. - Py_INCREF(ret.f2); } else { // offset is dummy in this case, just use the old value. ret.f1 = offset; @@ -393,7 +409,18 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) { Py_ssize_t py_offset = CPyTagged_AsSsize_t(offset); if (PyDict_CheckExact(dict_or_iter)) { - ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &ret.f2, &ret.f3); + PyObject *key; + PyObject *value; + // PyDict_Next() returns borrowed references. On free-threaded builds, + // hold the dict lock until we have converted them to strong references. + Py_BEGIN_CRITICAL_SECTION(dict_or_iter); + ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &key, &value); + if (ret.f0) { + ret.f2 = Py_NewRef(key); + ret.f3 = Py_NewRef(value); + } + Py_END_CRITICAL_SECTION(); + if (ret.f0) { ret.f1 = CPyTagged_FromSsize_t(py_offset); } else { @@ -401,6 +428,8 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) { ret.f1 = 0; ret.f2 = Py_None; ret.f3 = Py_None; + Py_INCREF(ret.f2); + Py_INCREF(ret.f3); } } else { ret.f1 = offset; @@ -413,6 +442,8 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) { ret.f0 = 0; ret.f2 = Py_None; ret.f3 = Py_None; + Py_INCREF(ret.f2); + Py_INCREF(ret.f3); } else { ret.f0 = 1; ret.f2 = PyTuple_GET_ITEM(item, 0); @@ -423,9 +454,6 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) { return ret; } } - // PyDict_Next() returns borrowed references. - Py_INCREF(ret.f2); - Py_INCREF(ret.f3); return ret; }