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)
-        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,6 +88,9 @@ def _cache_stub_file_map(version_info):
 
 def import_module_decorator(func):
     def wrapper(evaluator, import_names, parent_module_context, sys_path, prefer_stubs):
+        try:
+            actual_context_set = evaluator.module_cache.get(import_names)
+        except KeyError:
             if import_names == ('os', 'path'):
                 # This is a huge exception, we follow a nested import
                 # ``os.path``, because it's a very important one in Python
@@ -97,12 +100,16 @@ def import_module_decorator(func):
                     parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
                 actual_context_set = parent_module_context.py__getattribute__('path')
             else:
-            actual_context_set = func(
+                if parent_module_context is not None and parent_module_context.is_stub():
-                evaluator,
+                    parent_module_contexts = parent_module_context.non_stub_context_set
-                import_names,
+                else:
-                parent_module_context,
+                    parent_module_contexts = [parent_module_context]
-                sys_path,
+                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)