Fix conversion for contexts

jedi/api/__init__.py

@@ -39,7 +39,7 @@ from jedi.evaluate.syntax_tree import tree_name_to_contexts
 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.conversion import convert_names
+from jedi.evaluate.gradual.conversion import convert_names, convert_contexts
 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
@@ -260,15 +260,15 @@ class Script(object):
                 return []
 
         context = self._evaluator.create_context(self._get_module(), leaf)
-        definitions = helpers.evaluate_goto_definition(self._evaluator, context, leaf)
 
-        names = convert_names(
-            [s.name for s in definitions],
+        contexts = helpers.evaluate_goto_definition(self._evaluator, context, leaf)
+        contexts = convert_contexts(
+            contexts,
             only_stubs=only_stubs,
             prefer_stubs=prefer_stubs,
         )
 
-        defs = [classes.Definition(self._evaluator, name) for name in names]
+        defs = [classes.Definition(self._evaluator, c.name) for c in contexts]
         # The additional set here allows the definitions to become unique in an
         # API sense. In the internals we want to separate more things than in
         # the API.

jedi/api/classes.py

@@ -17,7 +17,7 @@ from jedi.evaluate import compiled
 from jedi.evaluate.imports import ImportName
 from jedi.evaluate.context import FunctionExecutionContext
 from jedi.evaluate.gradual.typeshed import StubModuleContext
-from jedi.evaluate.gradual.conversion import convert_names, \
+from jedi.evaluate.gradual.conversion import convert_names, convert_contexts, \
     stub_to_python_context_set
 from jedi.api.keywords import KeywordName
 
@@ -309,11 +309,12 @@ class BaseDefinition(object):
         if not self._name.is_context_name:
             return []
 
-        names = convert_names(
-            [c.name for c in self._name.infer()],
+        contexts = convert_contexts(
+            self._name.infer(),
             only_stubs=only_stubs,
             prefer_stubs=prefer_stubs,
         )
+        names = [c.name for c in contexts],
         return [self if n == self._name else Definition(self._evaluator, n)
                 for n in names]
 

jedi/evaluate/gradual/conversion.py

@@ -126,7 +126,8 @@ def _python_to_stub_names(names, fallback_to_python=False):
             yield name
 
 
-def convert_names(names, only_stubs, prefer_stubs):
+def convert_names(names, only_stubs=False, prefer_stubs=False):
+    assert not (only_stubs and prefer_stubs)
     with debug.increase_indent_cm('convert names'):
         if only_stubs or prefer_stubs:
             return _python_to_stub_names(names, fallback_to_python=prefer_stubs)
@@ -134,6 +135,23 @@ def convert_names(names, only_stubs, prefer_stubs):
             return _try_stub_to_python_names(names, prefer_stub_to_compiled=True)
 
 
+def convert_contexts(contexts, only_stubs=False, prefer_stubs=False):
+    assert not (only_stubs and prefer_stubs)
+    with debug.increase_indent_cm('convert contexts'):
+        if only_stubs or prefer_stubs:
+            return ContextSet.from_sets(
+                to_stub(context)
+                or (ContextSet({context}) if prefer_stubs else NO_CONTEXTS)
+                for context in contexts
+            )
+        else:
+            return ContextSet.from_sets(
+                stub_to_python_context_set(stub_context, ignore_compiled=True)
+                or ContextSet({stub_context})
+                for stub_context in contexts
+            )
+
+
 # TODO merge with _python_to_stub_names?
 def to_stub(context):
     if context.is_stub():