Refactor things so goto is working in both directions

jedi/api/__init__.py

@@ -39,6 +39,8 @@ from jedi.evaluate.context import ModuleContext
 from jedi.evaluate.base_context import ContextSet
 from jedi.evaluate.context.iterable import unpack_tuple_to_dict
 #from jedi.evaluate.gradual.typeshed import try_to_merge_with_stub
+from jedi.evaluate.gradual.stub_context import try_stubs_to_actual_context_set, \
+    try_stubs_to_actual_names
 from jedi.evaluate.gradual.utils import load_proper_stub_module
 
 # Jedi uses lots and lots of recursion. By setting this a little bit higher, we
@@ -252,6 +254,8 @@ class Script(object):
 
         context = self._evaluator.create_context(self._get_module(), leaf)
         definitions = helpers.evaluate_goto_definition(self._evaluator, context, leaf)
+        # We don't want stubs here we want the actual contexts, if possible.
+        definitions = try_stubs_to_actual_context_set(definitions)
 
         names = [s.name for s in definitions]
         defs = [classes.Definition(self._evaluator, name) for name in names]
@@ -304,6 +308,7 @@ class Script(object):
                 return isinstance(name, imports.SubModuleName)
 
         names = filter_follow_imports(names, check)
+        names = try_stubs_to_actual_names(names)
 
         defs = [classes.Definition(self._evaluator, d) for d in set(names)]
         return helpers.sorted_definitions(defs)

jedi/api/classes.py

@@ -16,6 +16,7 @@ from jedi.evaluate.imports import ImportName
 from jedi.evaluate.filters import ParamName
 from jedi.evaluate.context import FunctionExecutionContext, MethodContext
 from jedi.evaluate.gradual.typeshed import StubOnlyModuleContext
+from jedi.evaluate.gradual.stub_context import name_to_stub
 from jedi.api.keywords import KeywordName
 
 
@@ -310,8 +311,8 @@ class BaseDefinition(object):
             return [self]
 
         return [
-            Definition(self._evaluator, d.stub_context.name)
+            Definition(self._evaluator, stub_def.name)
-            for d in self._name.infer() if d.stub_context is not None
+            for stub_def in name_to_stub(self._name)
         ]
 
     def goto_assignments(self):

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()
+    stub_module = stub_context.get_root_context()
-    assert isinstance(stub_only_module, StubOnlyModuleContext), stub_only_module
+
-    non_stubs = stub_only_module.get_stub_contexts()
+    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:
+        evaluator.stub_module_cache[import_names] = stub
-            # Try to load the modules in a way where they are loaded
+        return stub
-            # 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
     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`.
     """

jedi/plugins/flask.py

@@ -7,19 +7,19 @@ class FlaskPlugin(BasePlugin):
         Handle "magic" Flask extension imports:
         ``flask.ext.foo`` is really ``flask_foo`` or ``flaskext.foo``.
         """
-        def wrapper(evaluator, import_names, module_context, sys_path):
+        def wrapper(evaluator, import_names, module_context, *args, **kwargs):
             if len(import_names) == 3 and import_names[:2] == ('flask', 'ext'):
                 # New style.
                 ipath = (u'flask_' + import_names[2]),
-                context_set = callback(evaluator, ipath, None, sys_path)
+                context_set = callback(evaluator, ipath, None, *args, **kwargs)
                 if context_set:
                     return context_set
-                context_set = callback(evaluator, (u'flaskext',), None, sys_path)
+                context_set = callback(evaluator, (u'flaskext',), None, *args, **kwargs)
                 return callback(
                     evaluator,
                     (u'flaskext', import_names[2]),
                     next(iter(context_set)),
-                    sys_path
+                    *args, **kwargs
                 )
-            return callback(evaluator, import_names, module_context, sys_path)
+            return callback(evaluator, import_names, module_context, *args, **kwargs)
         return wrapper