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():

test/run.py

@@ -126,8 +126,7 @@ from jedi.api.classes import Definition
 from jedi.api.completion import get_user_scope
 from jedi import parser_utils
 from jedi.api.environment import get_default_environment, get_system_environment
-from jedi.evaluate.gradual.conversion import stub_to_python_context_set
-from jedi.evaluate.base_context import ContextSet
+from jedi.evaluate.gradual.conversion import convert_contexts
 
 
 TEST_COMPLETIONS = 0
@@ -139,16 +138,6 @@ TEST_USAGES = 3
 grammar36 = parso.load_grammar(version='3.6')
 
 
-def try_stubs_to_python_context_set(stub_contexts, prefer_stub_to_compiled=False):
-    contexts = ContextSet.from_sets(
-        stub_to_python_context_set(stub_context, ignore_compiled=prefer_stub_to_compiled)
-        or ContextSet([stub_context])
-        for stub_context in stub_contexts
-    )
-    debug.dbg('Stubs to actual: %s to %s', stub_contexts, contexts)
-    return contexts
-
-
 class IntegrationTestCase(object):
     def __init__(self, test_type, correct, line_nr, column, start, line,
                  path=None, skip_version_info=None):
@@ -242,9 +231,8 @@ class IntegrationTestCase(object):
                 if user_context.api_type == 'function':
                     user_context = user_context.get_function_execution()
                 element.parent = user_context.tree_node
-                results = try_stubs_to_python_context_set(
+                results = convert_contexts(
                     evaluator.eval_element(user_context, element),
-                    prefer_stub_to_compiled=True
                 )
                 if not results:
                     raise Exception('Could not resolve %s on line %s'