Implement importlib.util.set_loader: a decorator to automatically set

brettcannon · brettcannon · commit 2cf03a820483 · 2009-03-10T05:17:37.000Z
__loader__ on modules.
diff --git a/Doc/library/importlib.rst b/Doc/library/importlib.rst
@@ -348,7 +348,15 @@ an :term:`importer`.
     loader should initialize as specified by :pep:`302`.
 
 
-.. function:: set_package(method)
+.. function:: set_loader(fxn)
+
+    A :term:`decorator` for a :term:`loader` to set the :attr:`__loader__`
+    attribute on loaded modules. If the attribute is already set the decorator
+    does nothing. It is assumed that the first positional argument to the
+    wrapped method is what :attr:`__loader__` should be set to.
+
+
+.. function:: set_package(fxn)
 
     A :term:`decorator` for a :term:`loader` to set the :attr:`__package__`
     attribute on the module returned by the loader. If :attr:`__package__` is
diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
@@ -1,10 +1,6 @@
 to do
 /////
 
-* Public API left to expose (w/ docs!)
-
-    + util.set_loader
-
 * Implement InspectLoader for BuiltinImporter and FrozenImporter.
 
     + Expose function to see if a frozen module is a package.
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
@@ -110,6 +110,17 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def set_loader(fxn):
+    """Set __loader__ on the returned module."""
+    def wrapper(self, *args, **kwargs):
+        module = fxn(self, *args, **kwargs)
+        if not hasattr(module, '__loader__'):
+            module.__loader__ = self
+        return module
+    wrap(wrapper, fxn)
+    return wrapper
+
+
 class BuiltinImporter:
 
     """Meta path loader for built-in modules.
@@ -132,6 +143,7 @@ def find_module(cls, fullname, path=None):
 
     @classmethod
     @set_package
+    @set_loader
     def load_module(cls, fullname):
         """Load a built-in module."""
         if fullname not in sys.builtin_module_names:
@@ -161,6 +173,7 @@ def find_module(cls, fullname, path=None):
 
     @classmethod
     @set_package
+    @set_loader
     def load_module(cls, fullname):
         """Load a frozen module."""
         if cls.find_module(fullname) is None:
@@ -249,13 +262,12 @@ def __init__(self, name, path, is_pkg):
 
     @check_name
     @set_package
+    @set_loader
     def load_module(self, fullname):
         """Load an extension module."""
         is_reload = fullname in sys.modules
         try:
-            module = imp.load_dynamic(fullname, self._path)
-            module.__loader__ = self
-            return module
+            return imp.load_dynamic(fullname, self._path)
         except:
             if not is_reload and fullname in sys.modules:
                 del sys.modules[fullname]
diff --git a/Lib/importlib/test/builtin/test_loader.py b/Lib/importlib/test/builtin/test_loader.py
@@ -15,7 +15,8 @@ class LoaderTests(abc.LoaderTests):
     assert 'errno' in sys.builtin_module_names
     name = 'errno'
 
-    verification = {'__name__': 'errno', '__package__': ''}
+    verification = {'__name__': 'errno', '__package__': '',
+                    '__loader__': machinery.BuiltinImporter}
 
     def verify(self, module):
         """Verify that the module matches against what it should have."""
diff --git a/Lib/importlib/test/extension/test_loader.py b/Lib/importlib/test/extension/test_loader.py
@@ -24,6 +24,8 @@ def test_module(self):
                                 ('__package__', '')]:
                 self.assertEqual(getattr(module, attr), value)
             self.assert_(ext_util.NAME in sys.modules)
+            self.assert_(isinstance(module.__loader__,
+                                    importlib._ExtensionFileLoader))
 
     def test_package(self):
         # Extensions are not found in packages.
diff --git a/Lib/importlib/test/frozen/test_loader.py b/Lib/importlib/test/frozen/test_loader.py
@@ -9,15 +9,16 @@ def test_module(self):
         with util.uncache('__hello__'):
             module = machinery.FrozenImporter.load_module('__hello__')
             check = {'__name__': '__hello__', '__file__': '<frozen>',
-                        '__package__': ''}
+                    '__package__': '', '__loader__': machinery.FrozenImporter}
             for attr, value in check.items():
                 self.assertEqual(getattr(module, attr), value)
 
     def test_package(self):
         with util.uncache('__phello__'):
             module = machinery.FrozenImporter.load_module('__phello__')
             check = {'__name__': '__phello__', '__file__': '<frozen>',
-                     '__package__': '__phello__', '__path__': ['__phello__']}
+                     '__package__': '__phello__', '__path__': ['__phello__'],
+                     '__loader__': machinery.FrozenImporter}
             for attr, value in check.items():
                 attr_value = getattr(module, attr)
                 self.assertEqual(attr_value, value,
@@ -28,7 +29,8 @@ def test_lacking_parent(self):
         with util.uncache('__phello__', '__phello__.spam'):
             module = machinery.FrozenImporter.load_module('__phello__.spam')
             check = {'__name__': '__phello__.spam', '__file__': '<frozen>',
-                     '__package__': '__phello__'}
+                    '__package__': '__phello__',
+                    '__loader__': machinery.FrozenImporter}
             for attr, value in check.items():
                 attr_value = getattr(module, attr)
                 self.assertEqual(attr_value, value,
diff --git a/Lib/importlib/util.py b/Lib/importlib/util.py
@@ -1,3 +1,4 @@
 """Utility code for constructing importers, etc."""
 from ._bootstrap import module_for_loader
+from ._bootstrap import set_loader
 from ._bootstrap import set_package
