Cleanup stub imports / caching

jedi/evaluate/__init__.py

@@ -149,15 +149,8 @@ class Evaluator(object):
                       sys_path=None, prefer_stubs=True):
         if sys_path is None:
             sys_path = self.get_sys_path()
-        try:
+        return self._import_module(import_names, parent_module_context,
-            return self.module_cache.get(import_names)
+                                   sys_path, prefer_stubs=prefer_stubs)
-        except KeyError:
-            pass
-
-        context_set = self._import_module(import_names, parent_module_context,
-                                          sys_path, prefer_stubs=prefer_stubs)
-        self.module_cache.add(import_names, context_set)
-        return context_set
 
     @property
     @evaluator_function_cache()

jedi/evaluate/gradual/typeshed.py

@@ -88,21 +88,28 @@ def _cache_stub_file_map(version_info):
 
 def import_module_decorator(func):
     def wrapper(evaluator, import_names, parent_module_context, sys_path, prefer_stubs):
-        if import_names == ('os', 'path'):
+        try:
-            # This is a huge exception, we follow a nested import
+            actual_context_set = evaluator.module_cache.get(import_names)
-            # ``os.path``, because it's a very important one in Python
+        except KeyError:
-            # that is being achieved by messing with ``sys.modules`` in
+            if import_names == ('os', 'path'):
-            # ``os``.
+                # This is a huge exception, we follow a nested import
-            if parent_module_context is None:
+                # ``os.path``, because it's a very important one in Python
-                parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
+                # that is being achieved by messing with ``sys.modules`` in
-            actual_context_set = parent_module_context.py__getattribute__('path')
+                # ``os``.
-        else:
+                if parent_module_context is None:
-            actual_context_set = func(
+                    parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
-                evaluator,
+                actual_context_set = parent_module_context.py__getattribute__('path')
-                import_names,
+            else:
-                parent_module_context,
+                if parent_module_context is not None and parent_module_context.is_stub():
-                sys_path,
+                    parent_module_contexts = parent_module_context.non_stub_context_set
-            )
+                else:
+                    parent_module_contexts = [parent_module_context]
+                actual_context_set = ContextSet.from_sets(
+                    func(evaluator, import_names, p, sys_path,)
+                    for p in parent_module_contexts
+                )
+            evaluator.module_cache.add(import_names, actual_context_set)
+
         if not prefer_stubs:
             return actual_context_set
 

jedi/evaluate/imports.py

@@ -324,6 +324,7 @@ class Importer(object):
             force_unicode(i.value if isinstance(i, tree.Name) else i)
             for i in self.import_path
         )
+        sys_path = self._sys_path_with_modifications()
 
         context_set = [None]
         for i, name in enumerate(self.import_path):
@@ -331,9 +332,8 @@ class Importer(object):
                 self._evaluator.import_module(
                     import_names[:i+1],
                     parent_module_context,
-                    self._sys_path_with_modifications(),
+                    sys_path
-                )
+                ) for parent_module_context in context_set
-                for parent_module_context in context_set
             ])
             if not context_set:
                 message = 'No module named ' + '.'.join(import_names)