Goto stubs if on definition

jedi/api/__init__.py

@@ -38,6 +38,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.typeshed import try_to_merge_with_stub
 
 # Jedi uses lots and lots of recursion. By setting this a little bit higher, we
 # can remove some "maximum recursion depth" errors.
@@ -166,6 +167,9 @@ class Script(object):
             code_lines=self._code_lines,
             is_package=is_package,
         )
+        module, = try_to_merge_with_stub(
+            self._evaluator, None, module.string_names, ContextSet([module])
+        )
         self._evaluator.module_cache.add(names, ContextSet([module]))
         return module
 

jedi/evaluate/__init__.py

@@ -84,6 +84,7 @@ 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.stub_context import with_stub_context_if_possible
 
 
 def _execute(context, arguments):
@@ -273,10 +274,15 @@ class Evaluator(object):
         def_ = name.get_definition(import_name_always=True)
         if def_ is not None:
             type_ = def_.type
-            if type_ == 'classdef':
+            is_classdef = type_ == 'classdef'
-                return [ClassContext(self, context, name.parent)]
+            if is_classdef or type_ == 'funcdef':
+                if is_classdef:
+                    c = ClassContext(self, context, name.parent)
+                else:
+                    c = FunctionContext.from_context(context, name.parent)
+                return with_stub_context_if_possible(c)
             elif type_ == 'funcdef':
-                return [FunctionContext.from_context(context, name.parent)]
+                return []
 
             if type_ == 'expr_stmt':
                 is_simple_name = name.parent.type not in ('power', 'trailer')

jedi/evaluate/base_context.py

@@ -172,6 +172,10 @@ class Context(HelperContextMixin, BaseContext):
         debug.warning("Not possible to return the stop iterations of %s", self)
         return NO_CONTEXTS
 
+    def get_qualified_names(self):
+        # Returns Optional[List[str]]
+        return None
+
 
 def iterate_contexts(contexts, contextualized_node=None, is_async=False):
     """

jedi/evaluate/context/function.py

@@ -19,7 +19,6 @@ from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext, \
 from jedi.evaluate.context import iterable
 from jedi import parser_utils
 from jedi.evaluate.parser_cache import get_yield_exprs
-from jedi.evaluate.gradual.annotation import infer_return_types
 
 
 class LambdaName(AbstractNameDefinition):
@@ -38,7 +37,24 @@ class LambdaName(AbstractNameDefinition):
         return ContextSet([self._lambda_context])
 
 
-class FunctionMixin(object):
+class FunctionAndClassMixin(object):
+    def get_qualified_names(self):
+        if self.parent_context.is_class():
+            n = self.parent_context.get_qualified_names()
+            if n is None:
+                # This means that the parent class lives within a function.
+                return None
+            return n + [self.py__name__()]
+        elif self.parent_context.is_module():
+            return [self.py__name__()]
+        else:
+            return None
+
+    def py__name__(self):
+        return self.name.string_name
+
+
+class FunctionMixin(FunctionAndClassMixin):
     api_type = u'function'
 
     def get_filters(self, search_global=False, until_position=None, origin_scope=None):
@@ -82,9 +98,6 @@ class FunctionMixin(object):
 
         return FunctionExecutionContext(self.evaluator, self.parent_context, self, arguments)
 
-    def py__name__(self):
-        return self.name.string_name
-
 
 class FunctionContext(use_metaclass(CachedMetaClass, FunctionMixin, TreeContext)):
     """
@@ -179,6 +192,7 @@ class FunctionExecutionContext(TreeContext):
             returns = get_yield_exprs(self.evaluator, funcdef)
         else:
             returns = funcdef.iter_return_stmts()
+            from jedi.evaluate.gradual.annotation import infer_return_types
             context_set = infer_return_types(self)
             if context_set:
                 # If there are annotations, prefer them over anything else.

jedi/evaluate/context/klass.py

@@ -49,6 +49,7 @@ from jedi.evaluate.filters import ParserTreeFilter, TreeNameDefinition, \
 from jedi.evaluate.arguments import unpack_arglist
 from jedi.evaluate.base_context import ContextSet, iterator_to_context_set, \
     TreeContext, NO_CONTEXTS
+from jedi.evaluate.context.function import FunctionAndClassMixin
 
 
 def apply_py__get__(context, instance, class_context):
@@ -118,7 +119,7 @@ class ClassFilter(ParserTreeFilter):
         return [name for name in names if self._access_possible(name)]
 
 
-class ClassMixin(object):
+class ClassMixin(FunctionAndClassMixin):
     def is_class(self):
         return True
 
@@ -133,9 +134,6 @@ class ClassMixin(object):
     def name(self):
         return ContextName(self, self.tree_node.name)
 
-    def py__name__(self):
-        return self.name.string_name
-
     def get_param_names(self):
         for context_ in self.py__getattribute__(u'__init__'):
             if context_.is_function():

jedi/evaluate/gradual/stub_context.py

@@ -1,7 +1,7 @@
 from jedi.cache import memoize_method
 from jedi.parser_utils import get_call_signature_for_any
 from jedi.evaluate.utils import safe_property
-from jedi.evaluate.base_context import ContextWrapper
+from jedi.evaluate.base_context import ContextWrapper, ContextSet
 from jedi.evaluate.context.function import FunctionMixin, FunctionContext, MethodContext
 from jedi.evaluate.context.klass import ClassMixin, ClassContext
 from jedi.evaluate.context.module import ModuleMixin, ModuleContext
@@ -208,37 +208,55 @@ class StubName(NameWrapper):
         self._stub_name = stub_name
 
     @memoize_method
-    @iterator_to_context_set
     def infer(self):
         stub_contexts = self._stub_name.infer()
         if not stub_contexts:
-            for c in self._wrapped_name.infer():
+            return self._wrapped_name.infer()
-                yield c
-            return
 
         typ = self._wrapped_name.tree_name.parent.type
+        # TODO is this if a performance optimization?
         if typ in ('classdef', 'funcdef'):
             actual_context, = self._wrapped_name.infer()
-            for stub_context in stub_contexts:
+            return _add_stub_if_possible(self.parent_context, actual_context, stub_contexts)
-                if isinstance(stub_context, MethodContext):
-                    assert isinstance(actual_context, MethodContext)
-                    cls = StubMethodContext
-                elif isinstance(stub_context, FunctionContext):
-                    cls = StubFunctionContext
-                elif isinstance(stub_context, StubOnlyClass):
-                    cls = StubClassContext
-                else:
-                    yield stub_context
-                    continue
-                yield cls.create_cached(
-                    actual_context.evaluator,
-                    self.parent_context,
-                    actual_context,
-                    stub_context,
-                )
         else:
-            for c in stub_contexts:
+            return stub_contexts
-                yield c
+
+
+@iterator_to_context_set
+def _add_stub_if_possible(parent_context, actual_context, stub_contexts):
+    for stub_context in stub_contexts:
+        if isinstance(stub_context, MethodContext):
+            assert isinstance(actual_context, MethodContext)
+            cls = StubMethodContext
+        elif isinstance(stub_context, FunctionContext):
+            cls = StubFunctionContext
+        elif isinstance(stub_context, StubOnlyClass):
+            cls = StubClassContext
+        else:
+            yield stub_context
+            continue
+        yield cls.create_cached(
+            actual_context.evaluator,
+            parent_context,
+            actual_context,
+            stub_context,
+        )
+
+
+def with_stub_context_if_possible(actual_context):
+    names = actual_context.get_qualified_names()
+    stub_module = actual_context.get_root_context().stub_context
+    if stub_module is None:
+        return ContextSet([actual_context])
+
+    stub_contexts = ContextSet([stub_module])
+    for name in names:
+        stub_contexts = stub_contexts.py__getattribute__(name)
+    return _add_stub_if_possible(
+        actual_context.parent_context,
+        actual_context,
+        stub_contexts,
+    )
 
 
 class CompiledStubName(NameWrapper):

jedi/evaluate/gradual/typeshed.py

@@ -136,44 +136,49 @@ def import_module_decorator(func):
             else:
                 raise
 
-        import_name = import_names[-1]
+        return try_to_merge_with_stub(evaluator, parent_module_context,
-        map_ = None
+                                      import_names, context_set)
-        if len(import_names) == 1:
-            map_ = _cache_stub_file_map(evaluator.grammar.version_info)
-        elif isinstance(parent_module_context, StubModuleContext):
-            if not parent_module_context.stub_context.is_package:
-                # Only if it's a package (= a folder) something can be
-                # imported.
-                return context_set
-            path = parent_module_context.stub_context.py__path__()
-            map_ = _merge_create_stub_map(path)
-
-        if map_ is not None:
-            path = map_.get(import_name)
-            if path is not None:
-                try:
-                    stub_module_node = _load_stub(evaluator, path)
-                except FileNotFoundError:
-                    # The file has since been removed after looking for it.
-                    # TODO maybe empty cache?
-                    pass
-                else:
-                    if import_names == ('typing',):
-                        module_cls = TypingModuleWrapper
-                    else:
-                        module_cls = StubOnlyModuleContext
-                    file_name = os.path.basename(path)
-                    stub_module_context = module_cls(
-                        context_set, evaluator, stub_module_node,
-                        path=path,
-                        string_names=import_names,
-                        # The code was loaded with latest_grammar, so use
-                        # that.
-                        code_lines=get_cached_code_lines(evaluator.latest_grammar, path),
-                        is_package=file_name == '__init__.pyi',
-                    )
-                    modules = _merge_modules(context_set, stub_module_context)
-                    return ContextSet(modules)
-        # If no stub is found, just return the default.
-        return context_set
     return wrapper
+
+
+def try_to_merge_with_stub(evaluator, parent_module_context, import_names, actual_context_set):
+    import_name = import_names[-1]
+    map_ = None
+    if len(import_names) == 1:
+        map_ = _cache_stub_file_map(evaluator.grammar.version_info)
+    elif isinstance(parent_module_context, StubModuleContext):
+        if not parent_module_context.stub_context.is_package:
+            # Only if it's a package (= a folder) something can be
+            # imported.
+            return actual_context_set
+        path = parent_module_context.stub_context.py__path__()
+        map_ = _merge_create_stub_map(path)
+
+    if map_ is not None:
+        path = map_.get(import_name)
+        if path is not None:
+            try:
+                stub_module_node = _load_stub(evaluator, path)
+            except FileNotFoundError:
+                # The file has since been removed after looking for it.
+                # TODO maybe empty cache?
+                pass
+            else:
+                if import_names == ('typing',):
+                    module_cls = TypingModuleWrapper
+                else:
+                    module_cls = StubOnlyModuleContext
+                file_name = os.path.basename(path)
+                stub_module_context = module_cls(
+                    actual_context_set, evaluator, stub_module_node,
+                    path=path,
+                    string_names=import_names,
+                    # The code was loaded with latest_grammar, so use
+                    # that.
+                    code_lines=get_cached_code_lines(evaluator.latest_grammar, path),
+                    is_package=file_name == '__init__.pyi',
+                )
+                modules = _merge_modules(actual_context_set, stub_module_context)
+                return ContextSet(modules)
+    # If no stub is found, just return the default.
+    return actual_context_set

test/test_evaluate/test_gradual/test_typeshed.py

@@ -187,24 +187,24 @@ def test_goto_stubs(Script):
 
 @pytest.mark.parametrize(
     'code', [
-        'import os; os.walk'
+        'import os; os.walk',
-        'from collections import Counter; Counter'
+        'from collections import Counter; Counter',
     ])
 def test_goto_stubs_on_itself(Script, code):
     """
     If goto_stubs is used on an identifier in e.g. the stdlib, we should goto
     the stub of it.
     """
-    s = Script()
+    s = Script(code)
-    os_module, = s.goto_definitions()
+    def_, = s.goto_definitions()
-    stub = os_module.goto_stubs()
+    stub, = def_.goto_stubs()
 
     script_on_source = Script(
-        path=os_module.module_path,
+        path=def_.module_path,
-        line=os_module.line,
+        line=def_.line,
-        column=os_module.column
+        column=def_.column
     )
-    definition, = script_on_source.goto_assignments()
+    definition, = script_on_source.goto_definitions()
     same_stub, = definition.goto_stubs()
     assert stub.module_path == same_stub.module_path
     assert stub.line == same_stub.line