Add only_stubs and prefer_stubs as parameters to goto_assignments/goto_definitions

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 try_stub_to_actual_names
+from jedi.evaluate.gradual.conversion import convert_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
@@ -235,7 +235,7 @@ class Script(object):
         debug.speed('completions end')
         return completions
 
-    def goto_definitions(self):
+    def goto_definitions(self, **kwargs):
         """
         Return the definitions of a the path under the cursor.  goto function!
         This follows complicated paths and returns the end, not the first
@@ -245,8 +245,14 @@ class Script(object):
         because Python itself is a dynamic language, which means depending on
         an option you can have two different versions of a function.
 
+        :param only_stubs: Only return stubs for this goto call.
+        :param prefer_stubs: Prefer stubs to Python obects for this goto call.
         :rtype: list of :class:`classes.Definition`
         """
+        with debug.increase_indent_cm('goto_definitions'):
+            return self._goto_definitions(**kwargs)
+
+    def _goto_definitions(self, only_stubs=False, prefer_stubs=False):
         leaf = self._module_node.get_name_of_position(self._pos)
         if leaf is None:
             leaf = self._module_node.get_leaf_for_position(self._pos)
@@ -256,25 +262,40 @@ class Script(object):
         context = self._evaluator.create_context(self._get_module(), leaf)
         definitions = helpers.evaluate_goto_definition(self._evaluator, context, leaf)
 
-        names = [s.name for s in definitions]
+        names = convert_names(
+            [s.name for s in definitions],
+            only_stubs=only_stubs,
+            prefer_stubs=prefer_stubs,
+        )
+
         defs = [classes.Definition(self._evaluator, name) for name in names]
         # 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.
         return helpers.sorted_definitions(set(defs))
 
-    def goto_assignments(self, follow_imports=False, follow_builtin_imports=False):
+    def goto_assignments(self, follow_imports=False, follow_builtin_imports=False, **kwargs):
         """
         Return the first definition found, while optionally following imports.
         Multiple objects may be returned, because Python itself is a
         dynamic language, which means depending on an option you can have two
         different versions of a function.
 
+        .. note:: It is deprecated to use follow_imports and follow_builtin_imports as
+            positional arguments. Will be a keyword argument in 0.16.0.
+
         :param follow_imports: The goto call will follow imports.
         :param follow_builtin_imports: If follow_imports is True will decide if
             it follow builtin imports.
+        :param only_stubs: Only return stubs for this goto call.
+        :param prefer_stubs: Prefer stubs to Python obects for this goto call.
         :rtype: list of :class:`classes.Definition`
         """
+        with debug.increase_indent_cm('goto_assignments'):
+            return self._goto_assignments(follow_imports, follow_builtin_imports, **kwargs)
+
+    def _goto_assignments(self, follow_imports, follow_builtin_imports,
+                          only_stubs=False, prefer_stubs=False):
         def filter_follow_imports(names, check):
             for name in names:
                 if check(name):
@@ -297,13 +318,17 @@ class Script(object):
         if tree_name is None:
             # Without a name we really just want to jump to the result e.g.
             # executed by `foo()`, if we the cursor is after `)`.
-            return self.goto_definitions()
+            return self.goto_definitions(only_stubs=only_stubs, prefer_stubs=prefer_stubs)
         context = self._evaluator.create_context(self._get_module(), tree_name)
         names = list(self._evaluator.goto(context, tree_name))
 
         if follow_imports:
             names = filter_follow_imports(names, lambda name: name.is_import())
-        names = try_stub_to_actual_names(names, prefer_stub_to_compiled=True)
+        names = convert_names(
+            names,
+            only_stubs=only_stubs,
+            prefer_stubs=prefer_stubs,
+        )
 
         defs = [classes.Definition(self._evaluator, d) for d in set(names)]
         return helpers.sorted_definitions(defs)

jedi/api/classes.py

@@ -17,8 +17,8 @@ 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 try_stub_to_actual_names, \
+from jedi.evaluate.gradual.conversion import convert_names, \
-    stub_to_actual_context_set, try_stubs_to_actual_context_set, actual_to_stub_names
+    stub_to_actual_context_set
 from jedi.api.keywords import KeywordName
 
 
@@ -291,12 +291,11 @@ class BaseDefinition(object):
         if not self._name.is_context_name:
             return []
 
-        names = self._name.goto()
+        names = convert_names(
-        if only_stubs or prefer_stubs:
+            self._name.goto(),
-            names = actual_to_stub_names(names, fallback_to_actual=prefer_stubs)
+            only_stubs=only_stubs,
-        else:
+            prefer_stubs=prefer_stubs,
-            names = try_stub_to_actual_names(names, prefer_stub_to_compiled=True)
+        )
-
         return [self if n == self._name else Definition(self._evaluator, n)
                 for n in names]
 
@@ -310,13 +309,13 @@ class BaseDefinition(object):
         if not self._name.is_context_name:
             return []
 
-        # Param names are special because they are not handled by
+        names = convert_names(
-        # the evaluator method.
+            [c.name for c in self._name.infer()],
-        context_set = try_stubs_to_actual_context_set(
+            only_stubs=only_stubs,
-            self._name.infer(),
+            prefer_stubs=prefer_stubs,
-            prefer_stub_to_compiled=True,
         )
-        return [Definition(self._evaluator, d.name) for d in context_set]
+        return [self if n == self._name else Definition(self._evaluator, n)
+                for n in names]
 
     @property
     @memoize_method

jedi/api/helpers.py

@@ -137,13 +137,10 @@ def get_stack_at_position(grammar, code_lines, module_node, pos):
     )
 
 
-def evaluate_goto_definition(evaluator, context, leaf, prefer_stubs=False):
+def evaluate_goto_definition(evaluator, context, leaf):
     if leaf.type == 'name':
         # In case of a name we can just use goto_definition which does all the
         # magic itself.
-        if prefer_stubs:
-            return evaluator.goto_stub_definitions(context, leaf)
-        else:
         return evaluator.goto_definitions(context, leaf)
 
     parent = leaf.parent
@@ -156,13 +153,7 @@ def evaluate_goto_definition(evaluator, context, leaf, prefer_stubs=False):
         return eval_atom(context, leaf)
     elif leaf.type in ('fstring_string', 'fstring_start', 'fstring_end'):
         return get_string_context_set(evaluator)
-    if prefer_stubs:
     return definitions
-    from jedi.evaluate.gradual.conversion import try_stubs_to_actual_context_set
-    return try_stubs_to_actual_context_set(
-        definitions,
-        prefer_stub_to_compiled=True,
-    )
 
 
 CallSignatureDetails = namedtuple(
@@ -268,5 +259,4 @@ def cache_call_signatures(evaluator, context, bracket_leaf, code_lines, user_pos
         evaluator,
         context,
         bracket_leaf.get_previous_leaf(),
-        prefer_stubs=True,
     )

jedi/evaluate/__init__.py

@@ -84,8 +84,6 @@ from jedi.evaluate.context import ClassContext, FunctionContext, \
 from jedi.evaluate.context.iterable import CompForContext
 from jedi.evaluate.syntax_tree import eval_trailer, eval_expr_stmt, \
     eval_node, check_tuple_assignments
-from jedi.evaluate.gradual.conversion import try_stub_to_actual_names, \
-    actual_to_stub_names, try_stubs_to_actual_context_set
 
 
 def _execute(context, arguments):
@@ -257,13 +255,6 @@ class Evaluator(object):
         return eval_node(context, element)
 
     def goto_definitions(self, context, name):
-        # We don't want stubs here we want the actual contexts, if possible.
-        return try_stubs_to_actual_context_set(
-            self.goto_stub_definitions(context, name),
-            prefer_stub_to_compiled=True
-        )
-
-    def goto_stub_definitions(self, context, name):
         def_ = name.get_definition(import_name_always=True)
         if def_ is not None:
             type_ = def_.type

jedi/evaluate/gradual/conversion.py

@@ -136,6 +136,14 @@ def actual_to_stub_names(names, fallback_to_actual=False):
             yield name
 
 
+def convert_names(names, only_stubs, prefer_stubs):
+    with debug.increase_indent_cm('convert names'):
+        if only_stubs or prefer_stubs:
+            return actual_to_stub_names(names, fallback_to_actual=prefer_stubs)
+        else:
+            return try_stub_to_actual_names(names, prefer_stub_to_compiled=True)
+
+
 def to_stub(context):
     if context.is_stub():
         return ContextSet([context])

test/test_evaluate/test_gradual/test_typeshed.py

@@ -3,9 +3,11 @@ import os
 import pytest
 from parso.utils import PythonVersionInfo
 
+from jedi.api.project import Project
 from jedi.evaluate.gradual import typeshed, stub_context
 from jedi.evaluate.context import TreeInstance, BoundMethod, FunctionContext, \
     MethodContext, ClassContext
+from test.helpers import root_dir
 
 TYPESHED_PYTHON3 = os.path.join(typeshed.TYPESHED_PATH, 'stdlib', '3')
 
@@ -234,3 +236,53 @@ def test_goto_stubs_on_itself(Script, code, type_):
 
     _assert_is_same(same_definition, definition)
     _assert_is_same(same_definition, same_definition2)
+
+
+@pytest.mark.parametrize('way', ['direct', 'indirect'])
+@pytest.mark.parametrize(
+    'kwargs', [
+        dict(only_stubs=False, prefer_stubs=False),
+        dict(only_stubs=False, prefer_stubs=True),
+        dict(only_stubs=True, prefer_stubs=False),
+    ]
+)
+@pytest.mark.parametrize(
+    ('code', 'full_name', 'has_stub', 'has_python'), [
+        ['import os; os.walk', 'os.walk', True, True],
+        ['from collections import Counter', 'collections.Counter', True, True],
+        ['from collections', 'collections', True, True],
+        ['from collections import Counter; Counter', 'collections.Counter', True, True],
+        ['from collections import Counter; Counter()', 'collections.Counter', True, True],
+        ['from collections import Counter; Counter.most_common',
+         'collections.Counter.most_common', True, True],
+
+        ['from keyword import kwlist; kwlist', 'typing.Sequence', True, True],
+        #['from keyword import kwlist', 'typing.Sequence', True, True],
+
+        ['import with_stub', 'with_stub', True, True],
+        ['import with_stub', 'with_stub', True, True],
+        ['import with_stub_folder.python_only', 'with_stub_folder.python_only', False, True],
+        ['import stub_only', 'stub_only', True, False],
+    ])
+def test_infer_and_goto(Script, code, full_name, has_stub, has_python, way, kwargs):
+    project = Project(os.path.join(root_dir, 'test', 'completion', 'stub_folder'))
+    s = Script(code, _project=project)
+    if way == 'direct':
+        defs = s.goto_definitions(**kwargs)
+    else:
+        goto_defs = s.goto_assignments()
+        defs = [d for goto_def in goto_defs for d in goto_def.infer(**kwargs)]
+
+    only_stubs = kwargs['only_stubs']
+    prefer_stubs = kwargs['prefer_stubs']
+    if not has_stub and only_stubs:
+        assert not defs
+    else:
+        assert defs
+
+    for d in defs:
+        if prefer_stubs and has_stub:
+            assert d.is_stub()
+        if only_stubs:
+            assert d.is_stub()
+        assert d.full_name == full_name