Make sys.lazy_modules match PEP and keep internal lazy submodules tracking internal

Author: DinoV

Include/internal/pycore_interp_structs.h

@@ -349,7 +349,14 @@ struct _import_state {
     int lazy_imports_mode;
     PyObject *lazy_imports_filter;
     PyObject *lazy_importing_modules;
+    // The set stored in sys.lazy_modules if values that have been
+    // lazily imported. This value is only for debugging/introspection
+    // purposes and is not used at runtime.
     PyObject *lazy_modules;
+    // Submodules that have been imported lazily from modules which have
+    // been imported lazily. When the module is imported we need to add
+    // a LazyImportObject which refers to the submodule on the module.
+    PyObject *lazy_pending_submodules;
 #ifdef Py_GIL_DISABLED
     PyMutex lazy_mutex;
 #endif

Lib/test/test_lazy_import/__init__.py

@@ -37,8 +37,7 @@ def test_basic_unused(self):
         """Lazy imported module should not be loaded if never accessed."""
         import test.test_lazy_import.data.basic_unused
         self.assertNotIn("test.test_lazy_import.data.basic2", sys.modules)
-        self.assertIn("test.test_lazy_import.data", sys.lazy_modules)
-        self.assertEqual(sys.lazy_modules["test.test_lazy_import.data"], {"basic2"})
+        self.assertIn("test.test_lazy_import.data.basic2", sys.lazy_modules)
 
     def test_sys_lazy_modules(self):
         try:
@@ -48,7 +47,7 @@ def test_sys_lazy_modules(self):
 
         self.assertFalse("test.test_lazy_import.data.basic2" in sys.modules)
         self.assertIn("test.test_lazy_import.data", sys.lazy_modules)
-        self.assertEqual(sys.lazy_modules["test.test_lazy_import.data"], {"basic2"})
+        self.assertIn("test.test_lazy_import.data.basic2", sys.lazy_modules)
         test.test_lazy_import.data.basic_from_unused.basic2
         self.assertNotIn("test.test_import.data", sys.lazy_modules)
 
@@ -569,8 +568,8 @@ def my_filter(name):
         self.assertIs(sys.get_lazy_imports_filter(), my_filter)
 
     def test_lazy_modules_attribute_is_dict(self):
-        """sys.lazy_modules should be a dict per PEP 810."""
-        self.assertIsInstance(sys.lazy_modules, dict)
+        """sys.lazy_modules should be a set per PEP 810."""
+        self.assertIsInstance(sys.lazy_modules, set)
 
     @support.requires_subprocess()
     def test_lazy_modules_tracks_lazy_imports(self):
@@ -579,8 +578,7 @@ def test_lazy_modules_tracks_lazy_imports(self):
             import sys
             initial_count = len(sys.lazy_modules)
             import test.test_lazy_import.data.basic_unused
-            assert "test.test_lazy_import.data" in sys.lazy_modules
-            assert sys.lazy_modules["test.test_lazy_import.data"] == {"basic2"}
+            assert "test.test_lazy_import.data.basic2" in sys.lazy_modules
             assert len(sys.lazy_modules) > initial_count
             print("OK")
         """)
@@ -1029,15 +1027,14 @@ def test_module_added_to_lazy_modules_on_lazy_import(self):
             lazy import test.test_lazy_import.data.basic2
 
             # Should be in lazy_modules after lazy import
-            assert "test.test_lazy_import.data" in sys.lazy_modules
-            assert sys.lazy_modules["test.test_lazy_import.data"] == {"basic2"}
+            assert "test.test_lazy_import.data.basic2" in sys.lazy_modules
             assert len(sys.lazy_modules) > initial_count
 
             # Trigger reification
             _ = test.test_lazy_import.data.basic2.x
 
             # Module should still be tracked (for diagnostics per PEP 810)
-            assert "test.test_lazy_import.data" not in sys.lazy_modules
+            assert "test.test_lazy_import.data.basic2" not in sys.lazy_modules
             print("OK")
         """)
         result = subprocess.run(
@@ -1050,8 +1047,8 @@ def test_module_added_to_lazy_modules_on_lazy_import(self):
 
     def test_lazy_modules_is_per_interpreter(self):
         """Each interpreter should have independent sys.lazy_modules."""
-        # Basic test that sys.lazy_modules exists and is a dict
-        self.assertIsInstance(sys.lazy_modules, dict)
+        # Basic test that sys.lazy_modules exists and is a set
+        self.assertIsInstance(sys.lazy_modules, set)
 
     def test_lazy_module_without_children_is_tracked(self):
         code = textwrap.dedent("""
@@ -1060,10 +1057,6 @@ def test_lazy_module_without_children_is_tracked(self):
             assert "json" in sys.lazy_modules, (
                 f"expected 'json' in sys.lazy_modules, got {set(sys.lazy_modules)}"
             )
-            assert sys.lazy_modules["json"] == set(), (
-                f"expected empty set for sys.lazy_modules['json'], "
-                f"got {sys.lazy_modules['json']!r}"
-            )
             print("OK")
         """)
         assert_python_ok("-c", code)
@@ -1932,7 +1925,7 @@ def create_lazy_imports(idx):
                 t.join()
 
             assert not errors, f"Errors: {errors}"
-            assert isinstance(sys.lazy_modules, dict), "sys.lazy_modules is not a dict"
+            assert isinstance(sys.lazy_modules, set), "sys.lazy_modules is not a dict"
             print("OK")
         """)
 

Lib/test/test_lazy_import/__main__.py

@@ -0,0 +1,3 @@
+import unittest
+
+unittest.main('test.test_lazy_import')

Python/import.c

@@ -94,6 +94,8 @@ static struct _inittab *inittab_copy = NULL;
     (interp)->imports.modules_by_index
 #define LAZY_MODULES(interp) \
     (interp)->imports.lazy_modules
+#define LAZY_PENDING_SUBMODULES(interp) \
+    (interp)->imports.lazy_pending_submodules
 #define IMPORTLIB(interp) \
     (interp)->imports.importlib
 #define OVERRIDE_MULTI_INTERP_EXTENSIONS_CHECK(interp) \
@@ -271,15 +273,19 @@ import_get_module(PyThreadState *tstate, PyObject *name)
 PyObject *
 _PyImport_InitLazyModules(PyInterpreterState *interp)
 {
-    assert(LAZY_MODULES(interp) == NULL);
-    LAZY_MODULES(interp) = PyDict_New();
+    assert(LAZY_MODULES(interp) == NULL &&
+           LAZY_PENDING_SUBMODULES(interp) == NULL);
+
+    LAZY_PENDING_SUBMODULES(interp) = PyDict_New();
+    LAZY_MODULES(interp) = PySet_New(0);
     return LAZY_MODULES(interp);
 }
 
 void
 _PyImport_ClearLazyModules(PyInterpreterState *interp)
 {
     Py_CLEAR(LAZY_MODULES(interp));
+    Py_CLEAR(LAZY_PENDING_SUBMODULES(interp));
 }
 
 static int
@@ -4339,7 +4345,7 @@ get_mod_dict(PyObject *module)
 // ensure we have the set for the parent module name in sys.lazy_modules.
 // Returns a new reference.
 static PyObject *
-ensure_lazy_submodules(PyDictObject *lazy_modules, PyObject *parent)
+ensure_lazy_pending_submodules(PyDictObject *lazy_modules, PyObject *parent)
 {
     PyObject *lazy_submodules;
     Py_BEGIN_CRITICAL_SECTION(lazy_modules);
@@ -4358,6 +4364,9 @@ ensure_lazy_submodules(PyDictObject *lazy_modules, PyObject *parent)
     return lazy_submodules;
 }
 
+// Ensures that we have a LazyImportObject on the parent module for
+// all children modules which have been lazily imported. If the parent
+// module overrides the child attribute then the value is not replaced.
 static int
 register_lazy_on_parent(PyThreadState *tstate, PyObject *name,
                         PyObject *builtins)
@@ -4369,16 +4378,16 @@ register_lazy_on_parent(PyThreadState *tstate, PyObject *name,
     PyObject *parent_dict = NULL;
 
     PyInterpreterState *interp = tstate->interp;
-    PyObject *lazy_modules = LAZY_MODULES(interp);
-    assert(lazy_modules != NULL);
+    PyObject *lazy_pending_submodules = LAZY_PENDING_SUBMODULES(interp);
+    assert(lazy_pending_submodules != NULL);
 
     Py_INCREF(name);
     while (true) {
         Py_ssize_t dot = PyUnicode_FindChar(name, '.', 0,
                                             PyUnicode_GET_LENGTH(name), -1);
         if (dot < 0) {
-            PyObject *lazy_submodules = ensure_lazy_submodules(
-                (PyDictObject *)lazy_modules, name);
+            PyObject *lazy_submodules = ensure_lazy_pending_submodules(
+                (PyDictObject *)lazy_pending_submodules, name);
             if (lazy_submodules == NULL) {
                 goto done;
             }
@@ -4400,8 +4409,8 @@ register_lazy_on_parent(PyThreadState *tstate, PyObject *name,
         }
 
         // Record the child as being lazily imported from the parent.
-        PyObject *lazy_submodules = ensure_lazy_submodules(
-            (PyDictObject *)lazy_modules, parent);
+        PyObject *lazy_submodules = ensure_lazy_pending_submodules(
+            (PyDictObject *)lazy_pending_submodules, parent);
         if (lazy_submodules == NULL) {
             goto done;
         }
@@ -4464,6 +4473,14 @@ register_from_lazy_on_parent(PyThreadState *tstate, PyObject *abs_name,
     if (fromname == NULL) {
         return -1;
     }
+
+    // Add the module name to sys.lazy_modules set (PEP 810).
+    PyObject *lazy_modules = LAZY_MODULES(tstate->interp);
+    if (PySet_Add(lazy_modules, fromname) < 0) {
+        Py_DECREF(fromname);
+        return -1;
+    }
+
     int res = register_lazy_on_parent(tstate, fromname, builtins);
     Py_DECREF(fromname);
     return res;
@@ -4555,6 +4572,13 @@ _PyImport_LazyImportModuleLevelObject(PyThreadState *tstate,
         Py_DECREF(abs_name);
         return NULL;
     }
+
+    // Add the module name to sys.lazy_modules set (PEP 810).
+    PyObject *lazy_modules = LAZY_MODULES(tstate->interp);
+    if (PySet_Add(lazy_modules, abs_name) < 0) {
+        goto error;
+    }
+
     if (fromlist && PyUnicode_Check(fromlist)) {
         if (register_from_lazy_on_parent(tstate, abs_name, fromlist,
                                          builtins) < 0) {
@@ -4791,6 +4815,7 @@ _PyImport_ClearCore(PyInterpreterState *interp)
     Py_CLEAR(IMPORTLIB(interp));
     Py_CLEAR(IMPORT_FUNC(interp));
     Py_CLEAR(LAZY_IMPORT_FUNC(interp));
+    Py_CLEAR(interp->imports.lazy_pending_submodules);
     Py_CLEAR(interp->imports.lazy_modules);
     Py_CLEAR(interp->imports.lazy_importing_modules);
     Py_CLEAR(interp->imports.lazy_imports_filter);
@@ -5636,11 +5661,13 @@ _imp__set_lazy_attributes_impl(PyObject *module, PyObject *modobj,
     PyThreadState *tstate = _PyThreadState_GET();
     PyObject *module_dict = NULL;
     PyObject *ret = NULL;
-    PyObject *lazy_modules = LAZY_MODULES(tstate->interp);
-    assert(lazy_modules != NULL);
+    PyObject *lazy_pending_modules = LAZY_PENDING_SUBMODULES(tstate->interp);
+    assert(lazy_pending_modules != NULL);
 
     PyObject *lazy_submodules;
-    if (PyDict_GetItemRef(lazy_modules, name, &lazy_submodules) < 0) {
+    if (PySet_Discard(LAZY_MODULES(tstate->interp), name) < 0) {
+        return NULL;
+    } else if (PyDict_GetItemRef(lazy_pending_modules, name, &lazy_submodules) < 0) {
         return NULL;
     }
     else if (lazy_submodules == NULL) {
@@ -5659,8 +5686,7 @@ _imp__set_lazy_attributes_impl(PyObject *module, PyObject *modobj,
     Py_END_CRITICAL_SECTION();
     Py_DECREF(lazy_submodules);
 
-    // once a module is imported it is removed from sys.lazy_modules
-    if (PyDict_DelItem(lazy_modules, name) < 0) {
+    if (PyDict_DelItem(lazy_pending_modules, name) < 0) {
         goto error;
     }