Goto definitions goto stubs now have a proper implementation

jedi/api/__init__.py

@@ -39,6 +39,7 @@ 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
+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
 # can remove some "maximum recursion depth" errors.
@@ -150,7 +151,7 @@ class Script(object):
     # be called multiple times.
     @cache.memoize_method
     def _get_module(self):
-        names = ('__main__',)
+        names = None
         is_package = False
         if self.path is not None:
             import_names, is_p = transform_path_to_dotted(
@@ -161,6 +162,20 @@ class Script(object):
                 names = import_names
                 is_package = is_p
 
+        if self.path is not None and self.path.endswith('.pyi'):
+            # We are in a stub file. Try to load the stub properly.
+            stub_module = load_proper_stub_module(
+                self._evaluator,
+                cast_path(self.path),
+                names,
+                self._module_node
+            )
+            if stub_module is not None:
+                return stub_module
+
+        if names is None:
+            names = ('__main__',)
+
         module = ModuleContext(
             self._evaluator, self._module_node, cast_path(self.path),
             string_names=names,

jedi/api/classes.py

@@ -321,7 +321,13 @@ class BaseDefinition(object):
         return [Definition(self._evaluator, n) for n in names]
 
     def infer(self):
-        return [Definition(self._evaluator, d.name) for d in self._name.infer()]
+        tree_name = self._name.tree_name
+        parent_context = self._name.parent_context
+        if tree_name is None or parent_context is None:
+            context_set = self._name.infer()
+        else:
+            context_set = self._evaluator.goto_definitions(parent_context, tree_name)
+        return [Definition(self._evaluator, d.name) for d in context_set]
 
     @property
     @memoize_method

jedi/evaluate/__init__.py

@@ -84,7 +84,8 @@ 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
+from jedi.evaluate.gradual.stub_context import with_stub_context_if_possible, \
+    stub_to_actual_context_set
 
 
 def _execute(context, arguments):
@@ -138,7 +139,8 @@ class Evaluator(object):
             self,
         )
 
-    def import_module(self, import_names, parent_module_context=None, sys_path=None):
+    def import_module(self, import_names, parent_module_context=None,
+                      sys_path=None, load_stub=True):
         if sys_path is None:
             sys_path = self.get_sys_path()
         try:
@@ -146,7 +148,8 @@ class Evaluator(object):
         except KeyError:
             pass
 
-        context_set = self._import_module(import_names, parent_module_context, sys_path)
+        context_set = self._import_module(import_names, parent_module_context,
+                                          sys_path, load_stub=load_stub)
         self.module_cache.add(import_names, context_set)
         return context_set
 
@@ -280,9 +283,10 @@ class Evaluator(object):
                     c = ClassContext(self, context, name.parent)
                 else:
                     c = FunctionContext.from_context(context, name.parent)
-                return with_stub_context_if_possible(c)
+                if context.is_stub():
-            elif type_ == 'funcdef':
+                    return stub_to_actual_context_set(c)
-                return []
+                else:
+                    return with_stub_context_if_possible(c)
 
             if type_ == 'expr_stmt':
                 is_simple_name = name.parent.type not in ('power', 'trailer')

jedi/evaluate/context/module.py

@@ -152,6 +152,14 @@ class ModuleContext(ModuleMixin, TreeContext):
         self.code_lines = code_lines
         self.is_package = is_package
 
+    def is_stub(self):
+        if self._path is not None and self._path.endswith('.pyi'):
+            # Currently this is the way how we identify stubs when e.g. goto is
+            # used in them. This could be changed if stubs would be identified
+            # sooner and used as StubOnlyModuleContext.
+            return True
+        return super(ModuleContext, self).is_stub()
+
     def py__name__(self):
         if self.string_names is None:
             return None
@@ -217,5 +225,5 @@ class ModuleContext(ModuleMixin, TreeContext):
         return "<%s: %s@%s-%s is_stub=%s>" % (
             self.__class__.__name__, self._string_name,
             self.tree_node.start_pos[0], self.tree_node.end_pos[0],
-            self._path is not None and self._path.endswith('.pyi')
+            self._path is not None and self.is_stub()
         )

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, ContextSet
+from jedi.evaluate.base_context import ContextWrapper, ContextSet, NO_CONTEXTS
 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
@@ -100,7 +100,7 @@ class _StubOnlyContextMixin(object):
     def get_stub_only_filter(self, parent_context, non_stub_filters, **filter_kwargs):
         # Here we remap the names from stubs to the actual module. This is
         # important if type inferences is needed in that module.
-        return StubFilter(
+        return _StubFilter(
             parent_context,
             non_stub_filters,
             self._get_stub_only_filters(**filter_kwargs),
@@ -111,7 +111,7 @@ class _StubOnlyContextMixin(object):
                           until_position=None, origin_scope=None):
         next(filters)  # Ignore the first filter and replace it with our own
         yield self.get_stub_only_filter(
-            parent_context=self,
+            parent_context=None,
             non_stub_filters=list(self._get_first_non_stub_filters()),
             search_global=search_global,
             until_position=until_position,
@@ -246,13 +246,13 @@ def _add_stub_if_possible(parent_context, actual_context, stub_contexts):
 
 def with_stub_context_if_possible(actual_context):
     assert actual_context.tree_node.type in ('classdef', 'funcdef')
-    names = actual_context.get_qualified_names()
+    qualified_names = actual_context.get_qualified_names()
     stub_module = actual_context.get_root_context().stub_context
-    if stub_module is None:
+    if stub_module is None or qualified_names is None:
         return ContextSet([actual_context])
 
     stub_contexts = ContextSet([stub_module])
-    for name in names:
+    for name in qualified_names:
         stub_contexts = stub_contexts.py__getattribute__(name)
     return _add_stub_if_possible(
         actual_context.parent_context,
@@ -261,6 +261,19 @@ def with_stub_context_if_possible(actual_context):
     )
 
 
+def stub_to_actual_context_set(stub_context):
+    qualified_names = stub_context.get_qualified_names()
+    if qualified_names is None:
+        return NO_CONTEXTS
+
+    stub_only_module = stub_context.get_root_context()
+    assert isinstance(stub_only_module, StubOnlyModuleContext), stub_only_module
+    non_stubs = stub_only_module.non_stub_context_set
+    for name in qualified_names:
+        non_stubs = non_stubs.py__getattribute__(name)
+    return non_stubs
+
+
 class CompiledStubName(NameWrapper):
     def __init__(self, parent_context, compiled_name, stub_name):
         super(CompiledStubName, self).__init__(stub_name)
@@ -333,12 +346,12 @@ class StubOnlyFilter(ParserTreeFilter):
         return True
 
 
-class StubFilter(AbstractFilter):
+class _StubFilter(AbstractFilter):
     """
     Merging names from stubs and non-stubs.
     """
     def __init__(self, parent_context, non_stub_filters, stub_filters, add_non_stubs):
-        self._parent_context = parent_context
+        self._parent_context = parent_context  # Optional[Context]
         self._non_stub_filters = non_stub_filters
         self._stub_filters = stub_filters
         self._add_non_stubs = add_non_stubs
@@ -393,9 +406,17 @@ class StubFilter(AbstractFilter):
                     stub_name = TypingModuleName(stub_name)
 
                 if isinstance(name, CompiledName):
-                    result.append(CompiledStubName(self._parent_context, name, stub_name))
+                    result.append(CompiledStubName(
+                        self._parent_context or stub_name.parent_context,
+                        name,
+                        stub_name
+                    ))
                 else:
-                    result.append(StubName(self._parent_context, name, stub_name))
+                    result.append(StubName(
+                        self._parent_context or name.parent_context,
+                        name,
+                        stub_name
+                    ))
         return result
 
     def __repr__(self):

jedi/evaluate/gradual/typeshed.py

@@ -11,7 +11,7 @@ from jedi.evaluate.gradual.stub_context import StubModuleContext, \
     TypingModuleWrapper, StubOnlyModuleContext
 
 _jedi_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-_TYPESHED_PATH = os.path.join(_jedi_path, 'third_party', 'typeshed')
+TYPESHED_PATH = os.path.join(_jedi_path, 'third_party', 'typeshed')
 
 
 def _merge_create_stub_map(directories):
@@ -50,7 +50,7 @@ def _create_stub_map(directory):
 def _get_typeshed_directories(version_info):
     check_version_list = ['2and3', str(version_info.major)]
     for base in ['stdlib', 'third_party']:
-        base = os.path.join(_TYPESHED_PATH, base)
+        base = os.path.join(TYPESHED_PATH, base)
         base_list = os.listdir(base)
         for base_list_entry in base_list:
             match = re.match(r'(\d+)\.(\d+)$', base_list_entry)
@@ -105,7 +105,7 @@ def _cache_stub_file_map(version_info):
 
 
 def import_module_decorator(func):
-    def wrapper(evaluator, import_names, parent_module_context, sys_path):
+    def wrapper(evaluator, import_names, parent_module_context, sys_path, load_stub=True):
         if import_names == ('_sqlite3',):
             # TODO Maybe find a better solution for this?
             # The problem is IMO how star imports are priorized and that
@@ -127,7 +127,7 @@ def import_module_decorator(func):
                 evaluator,
                 import_names,
                 parent_module_context,
-                sys_path
+                sys_path,
             )
         except JediImportError:
             if import_names == ('typing',):
@@ -136,6 +136,8 @@ def import_module_decorator(func):
             else:
                 raise
 
+        if not load_stub:
+            return context_set
         return try_to_merge_with_stub(evaluator, parent_module_context,
                                       import_names, context_set)
     return wrapper
@@ -164,21 +166,26 @@ def try_to_merge_with_stub(evaluator, parent_module_context, import_names, actua
                 # TODO maybe empty cache?
                 pass
             else:
-                if import_names == ('typing',):
+                return create_stub_module(evaluator, actual_context_set,
-                    module_cls = TypingModuleWrapper
+                                          stub_module_node, path, import_names)
-                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
+
+
+def create_stub_module(evaluator, actual_context_set, stub_module_node, path, import_names):
+    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)

jedi/evaluate/gradual/utils.py

@@ -0,0 +1,38 @@
+import os
+
+from jedi.evaluate.imports import JediImportError
+from jedi.evaluate.gradual.typeshed import TYPESHED_PATH, create_stub_module
+
+
+def load_proper_stub_module(evaluator, path, import_names, module_node):
+    """
+    This function is given a random .pyi file and should return the proper
+    module.
+    """
+    assert path.endswith('.pyi')
+    if path.startswith(TYPESHED_PATH):
+        # /foo/stdlib/3/os/__init__.pyi -> stdlib/3/os/__init__
+        rest = path[len(TYPESHED_PATH) + 1: -4]
+        split_paths = tuple(rest.split(os.path.sep))
+        # Remove the stdlib/3 or third_party/3.5 part
+        import_names = split_paths[2:]
+        if import_names[-1] == '__init__':
+            import_names = import_names[:-1]
+
+    if import_names is not None:
+        try:
+            actual_context_set = evaluator.import_module(import_names, load_stub=False)
+        except JediImportError as e:
+            return None
+
+        context_set = create_stub_module(
+            evaluator, actual_context_set, module_node, path, import_names
+        )
+        for m in context_set:
+            # Try to load the modules in a way where they are loaded
+            # correctly as stubs and not as actual modules (which is what
+            # will happen if this condition isn't True).
+            if m.stub_context.py__file__() == path:
+                evaluator.module_cache.add(import_names, context_set)
+                return m.stub_context
+    return None

jedi/evaluate/imports.py

@@ -438,7 +438,7 @@ class JediImportError(Exception):
 
 
 @import_module_decorator
-def import_module(evaluator, import_names, parent_module_context, sys_path):
+def import_module(evaluator, import_names, parent_module_context, sys_path, load_stub=True):
     """
     This method is very similar to importlib's `_gcd_import`.
     """

jedi/plugins/flask.py

@@ -8,12 +8,12 @@ class FlaskPlugin(BasePlugin):
         Handle "magic" Flask extension imports:
         ``flask.ext.foo`` is really ``flask_foo`` or ``flaskext.foo``.
         """
-        def wrapper(evaluator, import_names, module_context, sys_path):
+        def wrapper(evaluator, import_names, module_context, sys_path, load_stub):
             if len(import_names) == 3 and import_names[:2] == ('flask', 'ext'):
                 # New style.
                 ipath = (u'flask_' + import_names[2]),
                 try:
-                    return callback(evaluator, ipath, None, sys_path)
+                    return callback(evaluator, ipath, None, sys_path, load_stub)
                 except JediImportError:
                     context_set = callback(evaluator, (u'flaskext',), None, sys_path)
                     # If context_set has no content a JediImportError is raised
@@ -24,5 +24,5 @@ class FlaskPlugin(BasePlugin):
                         next(iter(context_set)),
                         sys_path
                     )
-            return callback(evaluator, import_names, module_context, sys_path)
+            return callback(evaluator, import_names, module_context, sys_path, load_stub)
         return wrapper

test/test_api/test_api.py

@@ -20,7 +20,7 @@ def test_preload_modules():
         # Filter the typeshed parser cache.
         typeshed_cache_count = sum(
             1 for path in grammar_cache
-            if path is not None and path.startswith(typeshed._TYPESHED_PATH)
+            if path is not None and path.startswith(typeshed.TYPESHED_PATH)
         )
         # +1 for None module (currently used)
         assert len(grammar_cache) - typeshed_cache_count == len(modules) + 1

test/test_evaluate/test_gradual/test_typeshed.py

@@ -7,13 +7,13 @@ from jedi.evaluate.gradual import typeshed, stub_context
 from jedi.evaluate.context import TreeInstance, BoundMethod, FunctionContext
 from jedi.evaluate.filters import TreeNameDefinition
 
-TYPESHED_PYTHON3 = os.path.join(typeshed._TYPESHED_PATH, 'stdlib', '3')
+TYPESHED_PYTHON3 = os.path.join(typeshed.TYPESHED_PATH, 'stdlib', '3')
 
 
 def test_get_typeshed_directories():
     def get_dirs(version_info):
         return {
-            d.replace(typeshed._TYPESHED_PATH, '').lstrip(os.path.sep)
+            d.replace(typeshed.TYPESHED_PATH, '').lstrip(os.path.sep)
             for d in typeshed._get_typeshed_directories(version_info)
         }
 
@@ -71,7 +71,7 @@ def test_keywords_variable(Script):
     def_, = Script(code).goto_definitions()
     assert def_.name == 'Sequence'
     # This points towards the typeshed implementation
-    assert typeshed._TYPESHED_PATH in def_.module_path
+    assert typeshed.TYPESHED_PATH in def_.module_path
 
 
 def test_class(Script):
@@ -134,7 +134,7 @@ def test_sys_hexversion(Script):
     def_, = script.completions()
     assert isinstance(def_._name, stub_context.CompiledStubName), def_._name
     assert isinstance(def_._name._wrapped_name, TreeNameDefinition)
-    assert typeshed._TYPESHED_PATH in def_.module_path
+    assert typeshed.TYPESHED_PATH in def_.module_path
     def_, = script.goto_definitions()
     assert def_.name == 'int'
 
@@ -204,30 +204,27 @@ def test_goto_stubs_on_itself(Script, code):
     """
     s = Script(code)
     def_, = s.goto_definitions()
-    #stub, = def_.goto_stubs()
+    stub, = def_.goto_stubs()
 
     script_on_source = Script(
         path=def_.module_path,
         line=def_.line,
         column=def_.column
     )
-    print('GO')
     definition, = script_on_source.goto_definitions()
-    print('\ta', definition._name._context, definition._name._context.parent_context)
-    return
     same_stub, = definition.goto_stubs()
     _assert_is_same(same_stub, stub)
     _assert_is_same(definition, def_)
     assert same_stub.module_path != def_.module_path
 
     # And the reverse.
-    script_on_source = Script(
+    script_on_stub = Script(
         path=same_stub.module_path,
         line=same_stub.line,
         column=same_stub.column
     )
 
-    same_definition, = script_on_source.goto_definitions()
+    same_definition, = script_on_stub.goto_definitions()
     same_definition2, = same_stub.infer()
     _assert_is_same(same_definition, definition)
     _assert_is_same(same_definition, same_definition2)