Goto stubs if on definition

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():
-                yield c
-            return
+            return self._wrapped_name.infer()
 
         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:
-                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,
-                )
+            return _add_stub_if_possible(self.parent_context, actual_context, stub_contexts)
         else:
-            for c in stub_contexts:
-                yield c
+            return stub_contexts
+
+
+@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]
-        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 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 try_to_merge_with_stub(evaluator, parent_module_context,
+                                      import_names, 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