Start adding tests for goto_assignments on stubs

jedi/evaluate/__init__.py

@@ -85,7 +85,7 @@ 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.stub_context import with_stub_context_if_possible, \
-    stub_to_actual_context_set
+    stub_to_actual_context_set, goto_with_stubs_if_possible
 
 
 def _execute(context, arguments):
@@ -342,7 +342,8 @@ class Evaluator(object):
             elif type_ == 'param':
                 return [ParamName(context, name)]
             elif type_ in ('funcdef', 'classdef'):
-                return [TreeNameDefinition(context, name)]
+                n = TreeNameDefinition(context, name)
+                return goto_with_stubs_if_possible(n)
             elif type_ in ('import_from', 'import_name'):
                 module_names = imports.infer_import(context, name, is_goto=True)
                 return module_names

jedi/evaluate/base_context.py

@@ -18,6 +18,13 @@ from jedi.evaluate.cache import evaluator_as_method_param_cache
 
 
 class HelperContextMixin(object):
+    def get_root_context(self):
+        context = self
+        while True:
+            if context.parent_context is None:
+                return context
+            context = context.parent_context
+
     @classmethod
     @evaluator_as_method_param_cache()
     def create_cached(cls, *args, **kwargs):

jedi/evaluate/context/module.py

@@ -136,6 +136,14 @@ class ModuleMixin(SubModuleDictMixin):
                 modules += new
         return modules
 
+    def get_qualified_names(self):
+        """
+        A module doesn't have a qualified name, but it's important to note that
+        it's reachable and not `None`. With this information we can add
+        qualified names on top for all context children.
+        """
+        return []
+
 
 class ModuleContext(ModuleMixin, TreeContext):
     api_type = u'module'

jedi/evaluate/gradual/stub_context.py

@@ -261,6 +261,27 @@ def with_stub_context_if_possible(actual_context):
     )
 
 
+def goto_with_stubs_if_possible(name):
+    root = name.parent_context.get_root_context()
+    stub = root.get_root_context().stub_context
+    if stub is None:
+        return [name]
+
+    qualified_names = name.parent_context.get_qualified_names()
+    if qualified_names is None:
+        return [name]
+
+    stub_contexts = ContextSet([stub])
+    for name in qualified_names:
+        stub_contexts = stub_contexts.py__getattribute__(name)
+    names = stub_contexts.py__getattribute__(name.string_name, is_goto=True)
+    return [
+        n
+        for n in names
+        if n.start_pos == name.start_pos and n.parent_context == name.parent_context
+    ] or [name]
+
+
 def stub_to_actual_context_set(stub_context):
     qualified_names = stub_context.get_qualified_names()
     if qualified_names is None:

test/test_evaluate/test_gradual/test_typeshed.py

@@ -192,18 +192,22 @@ def _assert_is_same(d1, d2):
     assert d1.column == d2.column
 
 
+@pytest.mark.parametrize('type_', ['goto', 'infer'])
 @pytest.mark.parametrize(
     'code', [
         'import os; os.walk',
         'from collections import Counter; Counter',
     ])
-def test_goto_stubs_on_itself(Script, code):
+def test_goto_stubs_on_itself(Script, code, type_):
     """
     If goto_stubs is used on an identifier in e.g. the stdlib, we should goto
     the stub of it.
     """
     s = Script(code)
-    def_, = s.goto_definitions()
+    if type_ == 'infer':
+        def_, = s.goto_definitions()
+    else:
+        def_, = s.goto_assignments(follow_imports=True)
     stub, = def_.goto_stubs()
 
     script_on_source = Script(
@@ -211,7 +215,10 @@ def test_goto_stubs_on_itself(Script, code):
         line=def_.line,
         column=def_.column
     )
-    definition, = script_on_source.goto_definitions()
+    if type_ == 'infer':
+        definition, = script_on_source.goto_definitions()
+    else:
+        definition, = script_on_source.goto_assignments()
     same_stub, = definition.goto_stubs()
     _assert_is_same(same_stub, stub)
     _assert_is_same(definition, def_)
@@ -224,7 +231,10 @@ def test_goto_stubs_on_itself(Script, code):
         column=same_stub.column
     )
 
-    same_definition, = script_on_stub.goto_definitions()
+    if type_ == 'infer':
+        same_definition, = script_on_stub.goto_definitions()
+    else:
+        same_definition, = script_on_stub.goto_assignments()
     same_definition2, = same_stub.infer()
     _assert_is_same(same_definition, definition)
     _assert_is_same(same_definition, same_definition2)