Issue #17117: Have both import itself and importlib.util.set_loader()

set __loader__ on a module when set to None.

Thanks to Gökcen Eraslan for the fix.

2 files changed, 76 insertions, 10 deletions

diff --git a/Lib/test/test_importlib/import_/test___loader__.py b/Lib/test/test_importlib/import_/test___loader__.py
new file mode 100644
index 0000000..0edbcc5
--- /dev/null
+++ b/Lib/test/test_importlib/import_/test___loader__.py
@@ -0,0 +1,44 @@
+import imp
+import sys
+import unittest
+
+from .. import util
+from . import util as import_util
+
+
+class LoaderMock:
+
+    def find_module(self, fullname, path=None):
+        return self
+
+    def load_module(self, fullname):
+        sys.modules[fullname] = self.module
+        return self.module
+
+
+class LoaderAttributeTests(unittest.TestCase):
+
+    def test___loader___missing(self):
+        module = imp.new_module('blah')
+        try:
+            del module.__loader__
+        except AttributeError:
+            pass
+        loader = LoaderMock()
+        loader.module = module
+        with util.uncache('blah'), util.import_state(meta_path=[loader]):
+            module = import_util.import_('blah')
+        self.assertEqual(loader, module.__loader__)
+
+    def test___loader___is_None(self):
+        module = imp.new_module('blah')
+        module.__loader__ = None
+        loader = LoaderMock()
+        loader.module = module
+        with util.uncache('blah'), util.import_state(meta_path=[loader]):
+            returned_module = import_util.import_('blah')
+        self.assertEqual(loader, module.__loader__)
+
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/Lib/test/test_importlib/test_util.py b/Lib/test/test_importlib/test_util.py
index efc8977..5f08719 100644
--- a/Lib/test/test_importlib/test_util.py
+++ b/Lib/test/test_importlib/test_util.py
@@ -162,6 +162,37 @@ class SetPackageTests(unittest.TestCase):
         self.assertEqual(wrapped.__qualname__, fxn.__qualname__)
 
 
+class SetLoaderTests(unittest.TestCase):
+
+    """Tests importlib.util.set_loader()."""
+
+    class DummyLoader:
+        @util.set_loader
+        def load_module(self, module):
+            return self.module
+
+    def test_no_attribute(self):
+        loader = self.DummyLoader()
+        loader.module = imp.new_module('blah')
+        try:
+            del loader.module.__loader__
+        except AttributeError:
+            pass
+        self.assertEqual(loader, loader.load_module('blah').__loader__)
+
+    def test_attribute_is_None(self):
+        loader = self.DummyLoader()
+        loader.module = imp.new_module('blah')
+        loader.module.__loader__ = None
+        self.assertEqual(loader, loader.load_module('blah').__loader__)
+
+    def test_not_reset(self):
+        loader = self.DummyLoader()
+        loader.module = imp.new_module('blah')
+        loader.module.__loader__ = 42
+        self.assertEqual(42, loader.load_module('blah').__loader__)
+
+
 class ResolveNameTests(unittest.TestCase):
 
     """Tests importlib.util.resolve_name()."""
@@ -195,14 +226,5 @@ class ResolveNameTests(unittest.TestCase):
             util.resolve_name('..bacon', 'spam')
 
 
-def test_main():
-    from test import support
-    support.run_unittest(
-            ModuleForLoaderTests,
-            SetPackageTests,
-            ResolveNameTests
-        )
-
-
 if __name__ == '__main__':
-    test_main()
+    unittest.main()