Add the flask plugin and move the import hacks there

jedi/evaluate/__init__.py

@@ -123,6 +123,10 @@ class Evaluator(object):
         from jedi.plugins import plugin_manager
         plugin_callbacks = plugin_manager.get_callbacks(self)
         self.execute = plugin_callbacks.decorate('execute', callback=execute)
+        self.import_module = plugin_callbacks.decorate(
+            'import_module',
+            callback=imports.import_module
+        )
 
     @property
     @evaluator_function_cache()

jedi/evaluate/imports.py

@@ -282,7 +282,7 @@ class Importer(object):
         if not self.import_path:
             return NO_CONTEXTS
 
-        return import_module(
+        return self._evaluator.import_module(
             self._evaluator,
             self.import_path,
             self.sys_path_with_modifications(),
@@ -387,16 +387,6 @@ def import_module(evaluator, import_path, sys_path, add_error_callback):
         for i in import_path
     ]
 
-    if len(import_path) > 2 and import_parts[:2] == ['flask', 'ext']:
-        # New style.
-        ipath = ('flask_' + str(import_parts[2]),) + import_path[3:]
-        modules = import_module(evaluator, ipath, sys_path, add_error_callback)
-        if modules:
-            return modules
-        else:
-            # Old style
-            return import_module(evaluator, ('flaskext',) + import_path[2:], sys_path, add_error_callback)
-
     if import_parts[0] in settings.auto_import_modules:
         module = _load_module(
             evaluator,
@@ -414,7 +404,7 @@ def import_module(evaluator, import_path, sys_path, add_error_callback):
     if len(import_path) > 1:
         # This is a recursive way of importing that works great with
         # the module cache.
-        bases = import_module(evaluator, import_path[:-1], sys_path, add_error_callback)
+        bases = evaluator.import_module(evaluator, import_path[:-1], sys_path, add_error_callback)
         if not bases:
             return NO_CONTEXTS
         # We can take the first element, because only the os special

jedi/plugins/__init__.py

@@ -1,5 +1,6 @@
 from jedi.plugins.stdlib import StdlibPlugin
 from jedi.plugins.typeshed import TypeshedPlugin
+from jedi.plugins.flask import FlaskPlugin
 
 
 class _PluginManager(object):
@@ -34,4 +35,5 @@ class _PluginCallbacks(object):
 plugin_manager = _PluginManager([
     StdlibPlugin,
     TypeshedPlugin,
+    FlaskPlugin,
 ])

jedi/plugins/base.py

@@ -12,7 +12,7 @@ class BasePlugin(object):
         """
         return callback
 
-    def import_module(callback):
+    def import_module(self, callback):
         """
         Decorates the
         import_module(evaluator, import_path, sys_path, add_error_callback)

jedi/plugins/flask.py

@@ -0,0 +1,29 @@
+from parso.python.tree import Name
+
+from jedi.plugins.base import BasePlugin
+
+
+class FlaskPlugin(BasePlugin):
+    def import_module(self, callback):
+        """
+        Handle "magic" Flask extension imports:
+        ``flask.ext.foo`` is really ``flask_foo`` or ``flaskext.foo``.
+        """
+        def wrapper(evaluator, import_path, *args, **kwargs):
+            import_parts = [
+                i.value if isinstance(i, Name) else i
+                for i in import_path
+            ]
+
+            if len(import_path) > 2 and import_parts[:2] == ['flask', 'ext']:
+                # New style.
+                ipath = ('flask_' + str(import_parts[2]),) + import_path[3:]
+                modules = callback(evaluator, ipath, *args, **kwargs)
+                if modules:
+                    return modules
+                else:
+                    # Old style
+                    return callback(evaluator, ('flaskext',) + import_path[2:], *args, **kwargs)
+            return callback(evaluator, import_path, *args, **kwargs)
+
+        return wrapper