diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py index 4ae422a1..4fea017c 100644 --- a/jedi/api/__init__.py +++ b/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 diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py index 9b22f70c..ae4eb985 100644 --- a/jedi/evaluate/__init__.py +++ b/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': - return [ClassContext(self, context, name.parent)] + is_classdef = type_ == 'classdef' + 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') diff --git a/jedi/evaluate/base_context.py b/jedi/evaluate/base_context.py index 4ee7aa35..7a510cf4 100644 --- a/jedi/evaluate/base_context.py +++ b/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): """ diff --git a/jedi/evaluate/context/function.py b/jedi/evaluate/context/function.py index 91affa41..e6fd7b66 100644 --- a/jedi/evaluate/context/function.py +++ b/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. diff --git a/jedi/evaluate/context/klass.py b/jedi/evaluate/context/klass.py index 2ca10f89..cab658c0 100644 --- a/jedi/evaluate/context/klass.py +++ b/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(): diff --git a/jedi/evaluate/gradual/stub_context.py b/jedi/evaluate/gradual/stub_context.py index e8e9408e..956c1029 100644 --- a/jedi/evaluate/gradual/stub_context.py +++ b/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): diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py index 8fd60c1d..0973a87f 100644 --- a/jedi/evaluate/gradual/typeshed.py +++ b/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 diff --git a/test/test_evaluate/test_gradual/test_typeshed.py b/test/test_evaluate/test_gradual/test_typeshed.py index d9019d06..306f3503 100644 --- a/test/test_evaluate/test_gradual/test_typeshed.py +++ b/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' - 'from collections import Counter; Counter' + 'import os; os.walk', + '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() - os_module, = s.goto_definitions() - stub = os_module.goto_stubs() + s = Script(code) + def_, = s.goto_definitions() + stub, = def_.goto_stubs() script_on_source = Script( - path=os_module.module_path, - line=os_module.line, - column=os_module.column + path=def_.module_path, + line=def_.line, + 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