gh-117182: Allow lazily loaded modules to modify their own __class__

3 files changed, 38 insertions, 4 deletions

diff --git a/Lib/importlib/util.py b/Lib/importlib/util.py
index f1bb4b1..c94a148 100644
--- a/Lib/importlib/util.py
+++ b/Lib/importlib/util.py
@@ -178,15 +178,17 @@ class _LazyModule(types.ModuleType):
             # Only the first thread to get the lock should trigger the load
             # and reset the module's class. The rest can now getattr().
             if object.__getattribute__(self, '__class__') is _LazyModule:
+                __class__ = loader_state['__class__']
+
                 # Reentrant calls from the same thread must be allowed to proceed without
                 # triggering the load again.
                 # exec_module() and self-referential imports are the primary ways this can
                 # happen, but in any case we must return something to avoid deadlock.
                 if loader_state['is_loading']:
-                    return object.__getattribute__(self, attr)
+                    return __class__.__getattribute__(self, attr)
                 loader_state['is_loading'] = True
 
-                __dict__ = object.__getattribute__(self, '__dict__')
+                __dict__ = __class__.__getattribute__(self, '__dict__')
 
                 # All module metadata must be gathered from __spec__ in order to avoid
                 # using mutated values.
@@ -216,8 +218,10 @@ class _LazyModule(types.ModuleType):
                 # Update after loading since that's what would happen in an eager
                 # loading situation.
                 __dict__.update(attrs_updated)
-                # Finally, stop triggering this method.
-                self.__class__ = types.ModuleType
+                # Finally, stop triggering this method, if the module did not
+                # already update its own __class__.
+                if isinstance(self, _LazyModule):
+                    object.__setattr__(self, '__class__', __class__)
 
         return getattr(self, attr)
 
diff --git a/Lib/test/test_importlib/test_lazy.py b/Lib/test/test_importlib/test_lazy.py
index 4d2cc4e..5c6e030 100644
--- a/Lib/test/test_importlib/test_lazy.py
+++ b/Lib/test/test_importlib/test_lazy.py
@@ -196,6 +196,34 @@ class LazyLoaderTests(unittest.TestCase):
             test_load = module.loads('{}')
             self.assertEqual(test_load, {})
 
+    def test_lazy_module_type_override(self):
+        # Verify that lazy loading works with a module that modifies
+        # its __class__ to be a custom type.
+
+        # Example module from PEP 726
+        module = self.new_module(source_code="""\
+import sys
+from types import ModuleType
+
+CONSTANT = 3.14
+
+class ImmutableModule(ModuleType):
+    def __setattr__(self, name, value):
+        raise AttributeError('Read-only attribute!')
+
+    def __delattr__(self, name):
+        raise AttributeError('Read-only attribute!')
+
+sys.modules[__name__].__class__ = ImmutableModule
+""")
+        sys.modules[TestingImporter.module_name] = module
+        self.assertIsInstance(module, util._LazyModule)
+        self.assertEqual(module.CONSTANT, 3.14)
+        with self.assertRaises(AttributeError):
+            module.CONSTANT = 2.71
+        with self.assertRaises(AttributeError):
+            del module.CONSTANT
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2024-03-23-12-28-05.gh-issue-117182.a0KANW.rst b/Misc/NEWS.d/next/Library/2024-03-23-12-28-05.gh-issue-117182.a0KANW.rst
new file mode 100644
index 0000000..6b3b841
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-03-23-12-28-05.gh-issue-117182.a0KANW.rst
@@ -0,0 +1,2 @@
+Lazy-loading of modules that modify their own ``__class__`` no longer
+reverts the ``__class__`` to :class:`types.ModuleType`.