From afced5014ccba872b30567e737efb2307477f694 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sun, 5 May 2019 22:52:48 +0200
Subject: [PATCH] Cleanup stub imports / caching

---
 jedi/evaluate/__init__.py         | 11 ++-------
 jedi/evaluate/gradual/typeshed.py | 37 ++++++++++++++++++-------------
 jedi/evaluate/imports.py          |  6 ++---
 3 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py
index 5d384aa2..db7db565 100644
--- a/jedi/evaluate/__init__.py
+++ b/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.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
+        return self._import_module(import_names, parent_module_context,
+                                   sys_path, prefer_stubs=prefer_stubs)
 
     @property
     @evaluator_function_cache()
diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py
index ad249465..976b923a 100644
--- a/jedi/evaluate/gradual/typeshed.py
+++ b/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'):
-            # 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 parent_module_context is None:
-                parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
-            actual_context_set = parent_module_context.py__getattribute__('path')
-        else:
-            actual_context_set = func(
-                evaluator,
-                import_names,
-                parent_module_context,
-                sys_path,
-            )
+        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
+                # that is being achieved by messing with ``sys.modules`` in
+                # ``os``.
+                if parent_module_context is None:
+                    parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
+                actual_context_set = parent_module_context.py__getattribute__('path')
+            else:
+                if parent_module_context is not None and parent_module_context.is_stub():
+                    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
 
diff --git a/jedi/evaluate/imports.py b/jedi/evaluate/imports.py
index 3ab4d7cc..bb2c5a34 100644
--- a/jedi/evaluate/imports.py
+++ b/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(),
-                )
-                for parent_module_context in context_set
+                    sys_path
+                ) for parent_module_context in context_set
             ])
             if not context_set:
                 message = 'No module named ' + '.'.join(import_names)