Refactor some of the import logic so it's possible to load typeshed modules

jedi/plugins/flask.py

@@ -10,20 +10,21 @@ class FlaskPlugin(BasePlugin):
         Handle "magic" Flask extension imports:
         ``flask.ext.foo`` is really ``flask_foo`` or ``flaskext.foo``.
         """
-        def wrapper(evaluator, import_names, *args, **kwargs):
-            if len(import_names) > 2 and import_names[:2] == ('flask', 'ext'):
+        def wrapper(evaluator, import_names, module_context, sys_path):
+            if len(import_names) == 3 and import_names[:2] == ('flask', 'ext'):
                 # New style.
-                ipath = ('flask_' + str(import_names[2]),) + import_names[3:]
+                ipath = ('flask_' + import_names[2]),
                 try:
-                    return callback(evaluator, ipath, *args, **kwargs)
+                    return callback(evaluator, ipath, None, sys_path)
                 except JediImportError:
-                    # Old style
+                    context_set = callback(evaluator, ('flaskext',), None, sys_path)
+                    # If context_set has no content a JediImportError is raised
+                    # which should be caught anyway by the caller.
                     return callback(
                         evaluator,
-                        ('flaskext',) + import_names[2:],
-                        *args,
-                        **kwargs
+                        ('flaskext', import_names[2]),
+                        next(iter(context_set)),
+                        sys_path
                     )
-            return callback(evaluator, import_names, *args, **kwargs)
-
+            return callback(evaluator, import_names, module_context, sys_path)
         return wrapper