Added tests for the import hook and fixed a problem with it.

flask/ext/__init__.py

@@ -8,7 +8,7 @@
     force all extensions to upgrade at the same time.
 
     When a user does ``from flask.ext.foo import bar`` it will attempt to
-    imprt ``from flask_foo import bar`` first and when that fails it will
+    import ``from flask_foo import bar`` first and when that fails it will
     try to import ``from flaskext.foo import bar``.
 
     We're switching from namespace packages because it was just too painful for
@@ -17,38 +17,46 @@
     :copyright: (c) 2011 by Armin Ronacher.
     :license: BSD, see LICENSE for more details.
 """
-import sys
 
 
 class _ExtensionImporter(object):
-    """This importer redirects imports from this submodule to other
-    locations.  This makes it possible to transition from the old
-    flaskext.name to the newer flask_name without people having a
-    hard time.
+    """This importer redirects imports from this submodule to other locations.
+    This makes it possible to transition from the old flaskext.name to the
+    newer flask_name without people having a hard time.
     """
     _module_choices = ['flask_%s', 'flaskext.%s']
-    _modules = sys.modules
+
+    def __init__(self):
+        from sys import meta_path
+        self.prefix = __name__ + '.'
+        self.prefix_cutoff = __name__.count('.') + 1
+
+        def _name(x):
+            cls = type(x)
+            return cls.__module__ + '.' + cls.__name__
+        this = _name(self)
+        meta_path[:] = [x for x in meta_path if _name(x) != this] + [self]
 
     def find_module(self, fullname, path=None):
-        if fullname.rsplit('.', 1)[0] == __name__:
+        if fullname.startswith(self.prefix):
             return self
 
     def load_module(self, fullname):
-        if fullname in self._modules:
-            return self._modules[fullname]
-        packname, modname = fullname.rsplit('.', 1)
+        from sys import modules
+        if fullname in modules:
+            return modules[fullname]
+        modname = fullname.split('.', self.prefix_cutoff)[self.prefix_cutoff]
         for path in self._module_choices:
             realname = path % modname
             try:
                 __import__(realname)
             except ImportError:
                 continue
-            module = self._modules[fullname] = self._modules[realname]
-            setattr(self._modules[__name__], modname, module)
-            module.__loader__ = self
+            module = modules[fullname] = modules[realname]
+            setattr(modules[__name__], modname, module)
             return module
         raise ImportError(fullname)
 
 
-sys.meta_path.append(_ExtensionImporter())
-del sys, _ExtensionImporter
+_ExtensionImporter()
+del _ExtensionImporter

flask/testsuite/ext.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""
+    flask.testsuite.ext
+    ~~~~~~~~~~~~~~~~~~~
+
+    Tests the extension import thing.
+
+    :copyright: (c) 2011 by Armin Ronacher.
+    :license: BSD, see LICENSE for more details.
+"""
+import sys
+import unittest
+from flask.testsuite import FlaskTestCase
+
+
+class ExtImportHookTestCase(FlaskTestCase):
+
+    def setup(self):
+        for entry, value in sys.modules.items():
+            if entry.startswith('flask.ext.') and value is not None:
+                sys.modules.pop(entry, None)
+        from flask import ext
+        reload(ext)
+
+        # reloading must not add more hooks
+        import_hooks = 0
+        for item in sys.meta_path:
+            cls = type(item)
+            if cls.__module__ == 'flask.ext' and \
+               cls.__name__ == '_ExtensionImporter':
+                import_hooks += 1
+        self.assert_equal(import_hooks, 1)
+
+    def test_flaskext_simple_import_normal(self):
+        from flask.ext.newext_simple import ext_id
+        self.assert_equal(ext_id, 'newext_simple')
+
+    def test_flaskext_simple_import_module(self):
+        from flask.ext import newext_simple
+        self.assert_equal(newext_simple.ext_id, 'newext_simple')
+        self.assert_equal(newext_simple.__name__, 'flask_newext_simple')
+
+    def test_flaskext_package_import_normal(self):
+        from flask.ext.newext_package import ext_id
+        self.assert_equal(ext_id, 'newext_package')
+
+    def test_flaskext_package_import_module(self):
+        from flask.ext import newext_package
+        self.assert_equal(newext_package.ext_id, 'newext_package')
+        self.assert_equal(newext_package.__name__, 'flask_newext_package')
+
+    def test_flaskext_package_import_submodule(self):
+        from flask.ext.newext_package import submodule
+        self.assert_equal(submodule.__name__, 'flask_newext_package.submodule')
+        self.assert_equal(submodule.test_function(), 42)
+
+
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(unittest.makeSuite(ExtImportHookTestCase))
+    return suite

flask/testsuite/test_apps/flask_newext_package/__init__.py

@@ -0,0 +1 @@
+ext_id = 'newext_package'

flask/testsuite/test_apps/flask_newext_package/submodule.py

@@ -0,0 +1,2 @@
+def test_function():
+    return 42

flask/testsuite/test_apps/flask_newext_simple.py

@@ -0,0 +1 @@
+ext_id = 'newext_simple'