From 27ab4ba3399a677e314836364b72df7803dc9f83 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Mon, 23 Jul 2018 04:04:07 +0200
Subject: [PATCH] Add the flask plugin and move the import hacks there

---
 jedi/evaluate/__init__.py |  4 ++++
 jedi/evaluate/imports.py  | 14 ++------------
 jedi/plugins/__init__.py  |  2 ++
 jedi/plugins/base.py      |  2 +-
 jedi/plugins/flask.py     | 29 +++++++++++++++++++++++++++++
 5 files changed, 38 insertions(+), 13 deletions(-)
 create mode 100644 jedi/plugins/flask.py

diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py
index 512e097a..067e1f32 100644
--- a/jedi/evaluate/__init__.py
+++ b/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()
diff --git a/jedi/evaluate/imports.py b/jedi/evaluate/imports.py
index 80bf1313..822e9bbc 100644
--- a/jedi/evaluate/imports.py
+++ b/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
diff --git a/jedi/plugins/__init__.py b/jedi/plugins/__init__.py
index 97318dc4..e69540e7 100644
--- a/jedi/plugins/__init__.py
+++ b/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,
 ])
diff --git a/jedi/plugins/base.py b/jedi/plugins/base.py
index 21033b57..370ed2a7 100644
--- a/jedi/plugins/base.py
+++ b/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)
diff --git a/jedi/plugins/flask.py b/jedi/plugins/flask.py
new file mode 100644
index 00000000..afdb9a85
--- /dev/null
+++ b/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