bpo-28411: Support other mappings in PyInterpreterState.modules. (#3593) · python/cpython@3f9eee6
@@ -1649,13 +1649,40 @@ getattribute(PyObject *obj, PyObject *name, int allow_qualname)
16491649return attr;
16501650}
165116511652+static int
1653+_checkmodule(PyObject *module_name, PyObject *module,
1654+PyObject *global, PyObject *dotted_path)
1655+{
1656+if (module == Py_None) {
1657+return -1;
1658+ }
1659+if (PyUnicode_Check(module_name) &&
1660+_PyUnicode_EqualToASCIIString(module_name, "__main__")) {
1661+return -1;
1662+ }
1663+1664+PyObject *candidate = get_deep_attribute(module, dotted_path, NULL);
1665+if (candidate == NULL) {
1666+if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
1667+PyErr_Clear();
1668+ }
1669+return -1;
1670+ }
1671+if (candidate != global) {
1672+Py_DECREF(candidate);
1673+return -1;
1674+ }
1675+Py_DECREF(candidate);
1676+return 0;
1677+}
1678+16521679static PyObject *
16531680whichmodule(PyObject *global, PyObject *dotted_path)
16541681{
16551682PyObject *module_name;
1656-PyObject *modules_dict;
1657-PyObject *module;
1683+PyObject *module = NULL;
16581684Py_ssize_t i;
1685+PyObject *modules;
16591686_Py_IDENTIFIER(__module__);
16601687_Py_IDENTIFIER(modules);
16611688_Py_IDENTIFIER(__main__);
@@ -1678,35 +1705,48 @@ whichmodule(PyObject *global, PyObject *dotted_path)
16781705assert(module_name == NULL);
1679170616801707/* Fallback on walking sys.modules */
1681-modules_dict = _PySys_GetObjectId(&PyId_modules);
1682-if (modules_dict == NULL) {
1708+modules = _PySys_GetObjectId(&PyId_modules);
1709+if (modules == NULL) {
16831710PyErr_SetString(PyExc_RuntimeError, "unable to get sys.modules");
16841711return NULL;
16851712 }
1686-1687-i = 0;
1688-while (PyDict_Next(modules_dict, &i, &module_name, &module)) {
1689-PyObject *candidate;
1690-if (PyUnicode_Check(module_name) &&
1691-_PyUnicode_EqualToASCIIString(module_name, "__main__"))
1692-continue;
1693-if (module == Py_None)
1694-continue;
1695-1696-candidate = get_deep_attribute(module, dotted_path, NULL);
1697-if (candidate == NULL) {
1698-if (!PyErr_ExceptionMatches(PyExc_AttributeError))
1713+if (PyDict_CheckExact(modules)) {
1714+i = 0;
1715+while (PyDict_Next(modules, &i, &module_name, &module)) {
1716+if (_checkmodule(module_name, module, global, dotted_path) == 0) {
1717+Py_INCREF(module_name);
1718+return module_name;
1719+ }
1720+if (PyErr_Occurred()) {
16991721return NULL;
1700-PyErr_Clear();
1701-continue;
1722+ }
17021723 }
1703-1704- if (candidate == global) {
1705- Py_INCREF(module_name);
1706- Py_DECREF(candidate);
1707-return module_name;
1724+ }
1725+else {
1726+PyObject *iterator = PyObject_GetIter(modules);
1727+if (iterator == NULL) {
1728+return NULL;
17081729 }
1709-Py_DECREF(candidate);
1730+while ((module_name = PyIter_Next(iterator))) {
1731+module = PyObject_GetItem(modules, module_name);
1732+if (module == NULL) {
1733+Py_DECREF(module_name);
1734+Py_DECREF(iterator);
1735+return NULL;
1736+ }
1737+if (_checkmodule(module_name, module, global, dotted_path) == 0) {
1738+Py_DECREF(module);
1739+Py_DECREF(iterator);
1740+return module_name;
1741+ }
1742+Py_DECREF(module);
1743+Py_DECREF(module_name);
1744+if (PyErr_Occurred()) {
1745+Py_DECREF(iterator);
1746+return NULL;
1747+ }
1748+ }
1749+Py_DECREF(iterator);
17101750 }
1711175117121752/* If no module is found, use __main__. */
@@ -6424,9 +6464,7 @@ _pickle_Unpickler_find_class_impl(UnpicklerObject *self,
64246464/*[clinic end generated code: output=becc08d7f9ed41e3 input=e2e6a865de093ef4]*/
64256465{
64266466PyObject *global;
6427-PyObject *modules_dict;
64286467PyObject *module;
6429-_Py_IDENTIFIER(modules);
6430646864316469/* Try to map the old names used in Python 2.x to the new ones used in
64326470 Python 3.x. We do this only with old pickle protocols and when the
@@ -6483,25 +6521,16 @@ _pickle_Unpickler_find_class_impl(UnpicklerObject *self,
64836521 }
64846522 }
648565236486-modules_dict = _PySys_GetObjectId(&PyId_modules);
6487-if (modules_dict == NULL) {
6488-PyErr_SetString(PyExc_RuntimeError, "unable to get sys.modules");
6489-return NULL;
6490- }
6491-6492-module = PyDict_GetItemWithError(modules_dict, module_name);
6524+module = PyImport_GetModule(module_name);
64936525if (module == NULL) {
64946526if (PyErr_Occurred())
64956527return NULL;
64966528module = PyImport_Import(module_name);
64976529if (module == NULL)
64986530return NULL;
6499-global = getattribute(module, global_name, self->proto >= 4);
6500-Py_DECREF(module);
6501- }
6502-else {
6503-global = getattribute(module, global_name, self->proto >= 4);
65046531 }
6532+global = getattribute(module, global_name, self->proto >= 4);
6533+Py_DECREF(module);
65056534return global;
65066535}
65076536