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

jedi/evaluate/imports.py

@@ -12,7 +12,6 @@ This module also supports import autocompletion, which means to complete
 statements like ``from datetim`` (cursor at the end would return ``datetime``).
 """
 import os
-from functools import partial
 
 from parso.python import tree
 from parso.tree import search_ancestor
@@ -287,20 +286,22 @@ class Importer(object):
             for i in self.import_path
         )
 
-        try:
-            return self._evaluator.import_module(
-                self._evaluator,
-                import_names,
-                self.sys_path_with_modifications(),
-
-            )
-        except JediImportError as e:
-            # Try to look up the name that was responsible for the import
-            # error. Since this is a plugin API, we intentionally just use
-            # strings, because every plugin would need to unpack the names.
-            name = self.import_path[len(e.import_names) - 1]
-            _add_error(self.module_context, name)
-            return NO_CONTEXTS
+        context_set = [None]
+        for i, name in enumerate(self.import_path):
+            try:
+                context_set = ContextSet.from_sets([
+                    self._evaluator.import_module(
+                        self._evaluator,
+                        import_names[:i+1],
+                        module_context,
+                        self.sys_path_with_modifications(),
+                    )
+                    for module_context in context_set
+                ])
+            except JediImportError:
+                _add_error(self.module_context, name)
+                return NO_CONTEXTS
+        return context_set
 
     def _generate_name(self, name, in_module=None):
         # Create a pseudo import to be able to follow them.
@@ -355,7 +356,8 @@ class Importer(object):
                 if context.api_type != 'module':  # not a module
                     continue
                 # namespace packages
-                if isinstance(context, ModuleContext) and context.py__file__().endswith('__init__.py'):
+                if isinstance(context, ModuleContext) \
+                        and context.py__file__().endswith('__init__.py'):
                     paths = context.py__path__()
                     names += self._get_module_names(paths, in_module=context)
 
@@ -395,7 +397,7 @@ class JediImportError(Exception):
         self.import_names = import_names
 
 
-def import_module(evaluator, import_names, sys_path):
+def import_module(evaluator, import_names, module_context, sys_path):
     """
     This method is very similar to importlib's `_gcd_import`.
     """
@@ -413,19 +415,21 @@ def import_module(evaluator, import_names, sys_path):
     except KeyError:
         pass
 
-    if len(import_names) > 1:
-        # This is a recursive way of importing that works great with
-        # the module cache.
-        bases = evaluator.import_module(evaluator, import_names[:-1], sys_path)
-        if not bases:
-            return NO_CONTEXTS
-        # We can take the first element, because only the os special
-        # case yields multiple modules, which is not important for
-        # further imports.
-        parent_module = list(bases)[0]
-
+    if module_context is None:
+        debug.dbg('global search_module %s', import_names[-1])
+        # Override the sys.path. It works only good that way.
+        # Injecting the path directly into `find_module` did not work.
+        code, module_path, is_pkg = evaluator.compiled_subprocess.get_module_info(
+            string=import_names[-1],
+            full_name=module_name,
+            sys_path=sys_path,
+            is_global_search=True,
+        )
+        if module_path is None:
+            raise JediImportError(import_names)
+    else:
         try:
-            method = parent_module.py__path__
+            method = module_context.py__path__
         except AttributeError:
             # The module is not a package.
             raise JediImportError(import_names)
@@ -447,18 +451,6 @@ def import_module(evaluator, import_names, sys_path):
                     break
             else:
                 raise JediImportError(import_names)
-    else:
-        debug.dbg('global search_module %s', import_names[-1])
-        # Override the sys.path. It works only good that way.
-        # Injecting the path directly into `find_module` did not work.
-        code, module_path, is_pkg = evaluator.compiled_subprocess.get_module_info(
-            string=import_names[-1],
-            full_name=module_name,
-            sys_path=sys_path,
-            is_global_search=True,
-        )
-        if module_path is None:
-            raise JediImportError(import_names)
 
     module = _load_module(
         evaluator, module_path, code, sys_path,

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

jedi/plugins/stdlib.py

@@ -78,14 +78,14 @@ class StdlibPlugin(BasePlugin):
         return wrapper
 
     def import_module(self, callback):
-        def wrapper(evaluator, import_names, sys_path):
+        def wrapper(evaluator, import_names, module_context, sys_path):
             # This is a huge exception, we follow a nested import
             # ``os.path``, because it's a very important one in Python
             # that is being achieved by messing with ``sys.modules`` in
             # ``os``.
             if import_names == ('os', 'path'):
-                return callback(evaluator, ('os',), sys_path).py__getattribute__('path')
-            return callback(evaluator, import_names, sys_path)
+                return module_context.py__getattribute__('path')
+            return callback(evaluator, import_names, module_context, sys_path)
 
 
         return wrapper