Refactor things so goto is working in both directions

jedi/evaluate/__init__.py

@@ -146,7 +146,7 @@ class Evaluator(object):
         )
 
     def import_module(self, import_names, parent_module_context=None,
-                      sys_path=None):
+                      sys_path=None, prefer_stubs=True):
         if sys_path is None:
             sys_path = self.get_sys_path()
         try:
@@ -154,7 +154,8 @@ class Evaluator(object):
         except KeyError:
             pass
 
-        context_set = self._import_module(import_names, parent_module_context, sys_path)
+        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
 

jedi/evaluate/gradual/stub_context.py

@@ -3,11 +3,11 @@ import os
 from jedi.cache import memoize_method
 from jedi.parser_utils import get_call_signature_for_any
 from jedi.evaluate.utils import safe_property
-from jedi.evaluate.base_context import ContextWrapper, ContextSet, NO_CONTEXTS
+from jedi.evaluate.base_context import ContextWrapper, ContextSet, \
+    NO_CONTEXTS, iterator_to_context_set
 from jedi.evaluate.context.function import FunctionMixin, FunctionContext, MethodContext
 from jedi.evaluate.context.klass import ClassMixin, ClassContext
 from jedi.evaluate.context.module import ModuleMixin, ModuleContext
-from jedi.evaluate.base_context import iterator_to_context_set
 from jedi.evaluate.filters import ParserTreeFilter, \
     NameWrapper, AbstractFilter, TreeNameDefinition
 from jedi.evaluate.compiled.context import CompiledName
@@ -326,14 +326,40 @@ def stub_to_actual_context_set(stub_context):
     if qualified_names is None:
         return NO_CONTEXTS
 
-    stub_only_module = stub_context.get_root_context()
-    assert isinstance(stub_only_module, StubOnlyModuleContext), stub_only_module
-    non_stubs = stub_only_module.get_stub_contexts()
+    stub_module = stub_context.get_root_context()
+
+    if not stub_module.is_stub():
+        return ContextSet([stub_context])
+
+    assert isinstance(stub_module, StubOnlyModuleContext), stub_module
+    non_stubs = stub_module.non_stub_context_set
     for name in qualified_names:
         non_stubs = non_stubs.py__getattribute__(name)
     return non_stubs
 
 
+def try_stubs_to_actual_context_set(stub_contexts):
+    return ContextSet.from_sets(
+        stub_to_actual_context_set(stub_context) or ContextSet(stub_context)
+        for stub_context in stub_contexts
+    )
+
+
+@to_list
+def try_stubs_to_actual_names(names):
+    for name in names:
+        parent_context = name.parent_context
+        if name.tree_name is None:
+            continue
+
+        if not parent_context.get_root_context().is_stub():
+            yield name
+            continue
+
+        for n in goto_non_stub(parent_context, name.tree_name):
+            yield n
+
+
 def stubify(parent_context, context):
     if parent_context.is_stub():
         return ContextSet(
@@ -374,6 +400,37 @@ def load_stubs(context):
     return stub_contexts
 
 
+def _load_stub_module(module):
+    if module.is_stub():
+        return module
+    from jedi.evaluate.gradual.typeshed import _try_to_load_stub
+    return _try_to_load_stub(
+        module.evaluator,
+        ContextSet([module]),
+        parent_module_context=None,
+        import_names=module.string_names
+    )
+
+
+def name_to_stub(name):
+    return ContextSet.from_sets(to_stub(c) for c in name.infer())
+
+
+def to_stub(context):
+    if context.is_stub():
+        return ContextSet([context])
+
+    qualified_names = context.get_qualified_names()
+    stub_module = _load_stub_module(context.get_root_context())
+    if stub_module is None or qualified_names is None:
+        return NO_CONTEXTS
+
+    stub_contexts = ContextSet([stub_module])
+    for name in qualified_names:
+        stub_contexts = stub_contexts.py__getattribute__(name)
+    return stub_contexts
+
+
 class CompiledStubName(NameWrapper):
     def __init__(self, parent_context, compiled_name, stub_name):
         super(CompiledStubName, self).__init__(stub_name)

jedi/evaluate/gradual/typeshed.py

@@ -88,14 +88,14 @@ def _cache_stub_file_map(version_info):
 
 
 def import_module_decorator(func):
-    def wrapper(evaluator, import_names, parent_module_context, sys_path):
+    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',))
+                parent_module_context, = evaluator.import_module(('os',), prefer_stubs=False)
             actual_context_set = parent_module_context.py__getattribute__('path')
         else:
             actual_context_set = func(
@@ -104,15 +104,23 @@ def import_module_decorator(func):
                 parent_module_context,
                 sys_path,
             )
+        if not prefer_stubs:
+            return actual_context_set
+
         stub = _try_to_load_stub(evaluator, actual_context_set, parent_module_context, import_names)
         if stub is not None:
-            return ContextSet(stub)
+            return ContextSet([stub])
         return actual_context_set
 
     return wrapper
 
 
 def _try_to_load_stub(evaluator, actual_context_set, parent_module_context, import_names):
+    try:
+        return evaluator.stub_module_cache[import_names]
+    except KeyError:
+        pass
+
     import_name = import_names[-1]
     map_ = None
     if len(import_names) == 1:
@@ -161,4 +169,4 @@ def create_stub_module(evaluator, actual_context_set, stub_module_node, path, im
         is_package=file_name == '__init__.pyi',
     )
     evaluator.stub_module_cache[import_names] = stub_module_context
-    return [stub_module_context]
+    return stub_module_context

jedi/evaluate/gradual/utils.py

@@ -19,18 +19,13 @@ def load_proper_stub_module(evaluator, path, import_names, module_node):
             import_names = import_names[:-1]
 
     if import_names is not None:
-        actual_context_set = evaluator.import_module(import_names)
+        actual_context_set = evaluator.import_module(import_names, prefer_stubs=False)
         if not actual_context_set:
             return None
 
-        context_set = create_stub_module(
+        stub = create_stub_module(
             evaluator, actual_context_set, module_node, path, import_names
         )
-        for m in context_set:
-            # Try to load the modules in a way where they are loaded
-            # correctly as stubs and not as actual modules (which is what
-            # will happen if this condition isn't True).
-            if m.stub_context.py__file__() == path:
-                evaluator.module_cache.add(import_names, context_set)
-                return m.stub_context
+        evaluator.stub_module_cache[import_names] = stub
+        return stub
     return None

jedi/evaluate/imports.py

@@ -424,7 +424,7 @@ class Importer(object):
 
 
 @import_module_decorator
-def import_module(evaluator, import_names, parent_module_context, sys_path, load_stub=True):
+def import_module(evaluator, import_names, parent_module_context, sys_path):
     """
     This method is very similar to importlib's `_gcd_import`.
     """