Reify lazy objects when accessed via the module object

Author: DinoV

Include/internal/pycore_moduleobject.h

@@ -24,6 +24,8 @@ typedef struct {
     PyObject *md_weaklist;
     // for logging purposes after md_dict is cleared
     PyObject *md_name;
+    // module version we last checked for lazy values
+    uint32_t m_dict_version;
 #ifdef Py_GIL_DISABLED
     void *md_gil;
 #endif

Lib/test/test_import/__init__.py

@@ -2629,7 +2629,30 @@ def test_compatibility_mode_relative(self):
             self.fail('lazy import failed')
 
         self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+    
+    def test_modules_dict(self):
+        try:
+            import test.test_import.data.lazy_imports.modules_dict
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_modules_geatattr(self):
+        try:
+            import test.test_import.data.lazy_imports.modules_getattr
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_modules_geatattr_other(self):
+        try:
+            import test.test_import.data.lazy_imports.modules_getattr_other
+        except ImportError as e:
+            self.fail('lazy import failed')
 
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
 
 
 class TestSinglePhaseSnapshot(ModuleSnapshot):

Lib/test/test_import/data/lazy_imports/modules_dict.py

@@ -0,0 +1,5 @@
+lazy import test.test_import.data.lazy_imports.basic2 as basic2
+
+import sys
+mod = sys.modules[__name__]
+x = mod.__dict__

Lib/test/test_import/data/lazy_imports/modules_getattr.py

@@ -0,0 +1,5 @@
+lazy import test.test_import.data.lazy_imports.basic2 as basic2
+
+import sys
+mod = sys.modules[__name__]
+x = mod.basic2

Lib/test/test_import/data/lazy_imports/modules_getattr_other.py

@@ -0,0 +1,5 @@
+lazy import test.test_import.data.lazy_imports.basic2 as basic2
+
+import sys
+mod = sys.modules[__name__]
+x = mod.__name__

Objects/moduleobject.c

@@ -6,6 +6,7 @@
 #include "pycore_fileutils.h"     // _Py_wgetcwd
 #include "pycore_import.h"        // _PyImport_GetNextModuleIndex()
 #include "pycore_interp.h"        // PyInterpreterState.importlib
+#include "pycore_lazyimportobject.h" // _PyLazyImportObject_Check()
 #include "pycore_long.h"          // _PyLong_GetOne()
 #include "pycore_modsupport.h"    // _PyModule_CreateInitialized()
 #include "pycore_moduleobject.h"  // _PyModule_GetDef()
@@ -22,12 +23,6 @@
     (assert(PyModule_Check(op)), _Py_CAST(PyModuleObject*, (op)))
 
 
-static PyMemberDef module_members[] = {
-    {"__dict__", _Py_T_OBJECT, offsetof(PyModuleObject, md_dict), Py_READONLY},
-    {0}
-};
-
-
 PyTypeObject PyModuleDef_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     "moduledef",                                /* tp_name */
@@ -1048,6 +1043,18 @@ _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
     PyObject *attr, *mod_name, *getattr;
     attr = _PyObject_GenericGetAttrWithDict((PyObject *)m, name, NULL, suppress);
     if (attr) {
+        if (PyLazyImport_CheckExact(attr)) {
+            PyObject *new_value = _PyImport_LoadLazyImportTstate(PyThreadState_GET(), attr);
+            if (new_value == NULL) {
+                return NULL;
+            } else if (PyDict_SetItem(m->md_dict, name, new_value) < 0) {
+                Py_DECREF(new_value);
+                Py_DECREF(attr);
+                return NULL;
+            }
+            Py_DECREF(attr);
+            return new_value;
+        }
         return attr;
     }
     if (suppress == 1) {
@@ -1273,7 +1280,7 @@ static PyMethodDef module_methods[] = {
 };
 
 static PyObject *
-module_get_dict(PyModuleObject *m)
+module_load_dict(PyModuleObject *m)
 {
     PyObject *dict = PyObject_GetAttr((PyObject *)m, &_Py_ID(__dict__));
     if (dict == NULL) {
@@ -1292,7 +1299,7 @@ module_get_annotate(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyModuleObject *m = _PyModule_CAST(self);
 
-    PyObject *dict = module_get_dict(m);
+    PyObject *dict = module_load_dict(m);
     if (dict == NULL) {
         return NULL;
     }
@@ -1317,7 +1324,7 @@ module_set_annotate(PyObject *self, PyObject *value, void *Py_UNUSED(ignored))
         return -1;
     }
 
-    PyObject *dict = module_get_dict(m);
+    PyObject *dict = module_load_dict(m);
     if (dict == NULL) {
         return -1;
     }
@@ -1347,7 +1354,7 @@ module_get_annotations(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyModuleObject *m = _PyModule_CAST(self);
 
-    PyObject *dict = module_get_dict(m);
+    PyObject *dict = module_load_dict(m);
     if (dict == NULL) {
         return NULL;
     }
@@ -1419,7 +1426,7 @@ module_set_annotations(PyObject *self, PyObject *value, void *Py_UNUSED(ignored)
 {
     PyModuleObject *m = _PyModule_CAST(self);
 
-    PyObject *dict = module_get_dict(m);
+    PyObject *dict = module_load_dict(m);
     if (dict == NULL) {
         return -1;
     }
@@ -1448,10 +1455,52 @@ module_set_annotations(PyObject *self, PyObject *value, void *Py_UNUSED(ignored)
     return ret;
 }
 
+static PyObject *
+module_get_dict(PyObject *mod, void *Py_UNUSED(ignored))
+{
+    PyModuleObject *self = (PyModuleObject *)mod;
+    PyThreadState *tstate = PyThreadState_GET();
+    Py_BEGIN_CRITICAL_SECTION(self->md_dict);
+    uint32_t version = _PyDict_GetKeysVersionForCurrentState(tstate->interp,
+                                                             (PyDictObject *)self->md_dict);
+    // Check if the dict has been updated since we last checked to see if
+    // it has lazy values.
+    if (self->m_dict_version != version || version == 0) {
+        // Scan for lazy values...
+        bool retry;
+        do {
+            retry = false;
+            Py_ssize_t pos = 0;
+            PyObject *key, *value;
+            while (PyDict_Next(self->md_dict, &pos, &key, &value)) {
+                if (PyLazyImport_CheckExact(value)) {
+                    PyObject *new_value = _PyImport_LoadLazyImportTstate(tstate, value);
+                    if (new_value == NULL) {
+                        return NULL;
+                    }
+                    if (PyDict_SetItem(self->md_dict, key, new_value) < 0) {
+                        Py_DECREF(new_value);
+                        return NULL;
+                    }
+                    if (!PyLazyImport_CheckExact(value)) {
+                        // Only force a retry if we actually made forward progress
+                        retry = true;
+                    }
+                    Py_DECREF(new_value);
+                }
+            }
+        } while(retry);
+        self->m_dict_version = _PyDict_GetKeysVersionForCurrentState(tstate->interp,
+                                                                     (PyDictObject *)self->md_dict);
+    }
+    Py_END_CRITICAL_SECTION();
+    return Py_NewRef(self->md_dict);
+}
 
 static PyGetSetDef module_getsets[] = {
     {"__annotations__", module_get_annotations, module_set_annotations},
     {"__annotate__", module_get_annotate, module_set_annotate},
+    {"__dict__", (getter)module_get_dict, NULL},
     {NULL}
 };
 
@@ -1485,7 +1534,7 @@ PyTypeObject PyModule_Type = {
     0,                                          /* tp_iter */
     0,                                          /* tp_iternext */
     module_methods,                             /* tp_methods */
-    module_members,                             /* tp_members */
+    0,                                          /* tp_members */
     module_getsets,                             /* tp_getset */
     0,                                          /* tp_base */
     0,                                          /* tp_dict */