diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py index bba7923b..e1e7d29d 100644 --- a/jedi/evaluate/__init__.py +++ b/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 diff --git a/jedi/evaluate/base_context.py b/jedi/evaluate/base_context.py index 7a510cf4..ed645ea8 100644 --- a/jedi/evaluate/base_context.py +++ b/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): diff --git a/jedi/evaluate/context/module.py b/jedi/evaluate/context/module.py index 0437e65c..b0c85dbd 100644 --- a/jedi/evaluate/context/module.py +++ b/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' diff --git a/jedi/evaluate/gradual/stub_context.py b/jedi/evaluate/gradual/stub_context.py index 5d144192..99c6b610 100644 --- a/jedi/evaluate/gradual/stub_context.py +++ b/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: diff --git a/test/test_evaluate/test_gradual/test_typeshed.py b/test/test_evaluate/test_gradual/test_typeshed.py index 11e13268..3bf2b48f 100644 --- a/test/test_evaluate/test_gradual/test_typeshed.py +++ b/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)