Expose importlib.util.set___package__.

3 files changed, 56 insertions, 22 deletions

diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
index 22103a1..72b7da8 100644
--- a/Lib/importlib/NOTES
+++ b/Lib/importlib/NOTES
@@ -1,10 +1,6 @@
 to do
 /////
 
-* Implement PEP 302 protocol for loaders (should just be a matter of testing).
-
-    + Source/bytecode.
-
 * Public API left to expose (w/ docs!)
 
     + abc
@@ -27,27 +23,15 @@ to do
             * get_code
             * get_source
 
-        - (?) SourceLoader(ResourceLoader)
+        - PyLoader(ResourceLoader)
 
             * source_path
-            * bytecode_path
-            * write_bytecode (not abstract)
-
-    + util
-
-        - set___package__ decorator
 
-    + machinery
+        - PyPycLoader(PyLoader)
 
-        - Extensions importers
-
-            * ExtensionFinder
-            * (?) Loader
-
-        - Source/bytecode importers
-
-            * SourceFinder
-            * (?) Loader
+            * source_mtime
+            * bytecode_path
+            * write_bytecode
 
     + test (Really want to worry about compatibility with future versions?)
 
diff --git a/Lib/importlib/test/test_util.py b/Lib/importlib/test/test_util.py
index 476b43f..8bd35f1 100644
--- a/Lib/importlib/test/test_util.py
+++ b/Lib/importlib/test/test_util.py
@@ -60,9 +60,58 @@ class ModuleForLoaderTests(unittest.TestCase):
             self.assert_(sys.modules[name] is module)
 
 
+class SetPackageTests(unittest.TestCase):
+
+
+    """Tests for importlib.util.set___package__."""
+
+    def verify(self, module, expect):
+        """Verify the module has the expected value for __package__ after
+        passing through set___package__."""
+        fxn = lambda: module
+        wrapped = util.set___package__(fxn)
+        wrapped()
+        self.assert_(hasattr(module, '__package__'))
+        self.assertEqual(expect, module.__package__)
+
+    def test_top_level(self):
+        # __package__ should be set to the empty string if a top-level module.
+        # Implicitly tests when package is set to None.
+        module = imp.new_module('module')
+        module.__package__ = None
+        self.verify(module, '')
+
+    def test_package(self):
+        # Test setting __package__ for a package.
+        module = imp.new_module('pkg')
+        module.__path__ = ['<path>']
+        module.__package__ = None
+        self.verify(module, 'pkg')
+
+    def test_submodule(self):
+        # Test __package__ for a module in a package.
+        module = imp.new_module('pkg.mod')
+        module.__package__ = None
+        self.verify(module, 'pkg')
+
+    def test_setting_if_missing(self):
+        # __package__ should be set if it is missing.
+        module = imp.new_module('mod')
+        if hasattr(module, '__package__'):
+            delattr(module, '__package__')
+        self.verify(module, '')
+
+    def test_leaving_alone(self):
+        # If __package__ is set and not None then leave it alone.
+        for value in (True, False):
+            module = imp.new_module('mod')
+            module.__package__ = value
+            self.verify(module, value)
+
+
 def test_main():
     from test import support
-    support.run_unittest(ModuleForLoaderTests)
+    support.run_unittest(ModuleForLoaderTests, SetPackageTests)
 
 
 if __name__ == '__main__':
diff --git a/Lib/importlib/util.py b/Lib/importlib/util.py
index 6250d5b..2b6154b 100644
--- a/Lib/importlib/util.py
+++ b/Lib/importlib/util.py
@@ -1,2 +1,3 @@
 """Utility code for constructing importers, etc."""
 from ._bootstrap import module_for_loader
+from ._bootstrap import set___package__