Add the flask plugin and move the import hacks there

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