diff --git a/jedi/inference/value/function.py b/jedi/inference/value/function.py index 7513c6a9..6f4da364 100644 --- a/jedi/inference/value/function.py +++ b/jedi/inference/value/function.py @@ -202,7 +202,7 @@ class MethodValue(FunctionValue): class BaseFunctionExecutionContext(ValueContext, TreeContextMixin): - def _infer_annotations(self): + def infer_annotations(self): raise NotImplementedError @inference_state_method_cache(default=NO_VALUES) @@ -216,7 +216,7 @@ class BaseFunctionExecutionContext(ValueContext, TreeContextMixin): value_set = NO_VALUES returns = get_yield_exprs(self.inference_state, funcdef) else: - value_set = self._infer_annotations() + value_set = self.infer_annotations() if value_set: # If there are annotations, prefer them over anything else. # This will make it faster. @@ -373,7 +373,7 @@ class FunctionExecutionContext(BaseFunctionExecutionContext): arguments=self._arguments ) - def _infer_annotations(self): + def infer_annotations(self): from jedi.inference.gradual.annotation import infer_return_types return infer_return_types(self._value, self._arguments) @@ -385,7 +385,7 @@ class FunctionExecutionContext(BaseFunctionExecutionContext): class AnonymousFunctionExecution(BaseFunctionExecutionContext): - def _infer_annotations(self): + def infer_annotations(self): # I don't think inferring anonymous executions is a big thing. # Anonymous contexts are mostly there for the user to work in. ~ dave return NO_VALUES diff --git a/jedi/inference/value/iterable.py b/jedi/inference/value/iterable.py index 8a60d31d..d2411a80 100644 --- a/jedi/inference/value/iterable.py +++ b/jedi/inference/value/iterable.py @@ -49,13 +49,12 @@ class GeneratorBase(LazyAttributeOverwrite, IterableMixin): array_type = None def _get_wrapped_value(self): - generator, = self.inference_state.typing_module \ - .py__getattribute__('Generator') \ - .execute_annotation() - return generator + instance, = self._get_cls().execute_annotation() + return instance - def is_instance(self): - return False + def _get_cls(self): + generator, = self.inference_state.typing_module.py__getattribute__('Generator') + return generator def py__bool__(self): return True @@ -77,6 +76,12 @@ class GeneratorBase(LazyAttributeOverwrite, IterableMixin): def name(self): return compiled.CompiledValueName(self, 'Generator') + def get_annotated_class_object(self): + from jedi.inference.gradual.generics import TupleGenericManager + gen_values = self.merge_types_of_iterate().py__class__() + gm = TupleGenericManager((gen_values, NO_VALUES, NO_VALUES)) + return self._get_cls().with_generics(gm) + class Generator(GeneratorBase): """Handling of `yield` functions.""" @@ -85,6 +90,9 @@ class Generator(GeneratorBase): self._func_execution_context = func_execution_context def py__iter__(self, contextualized_node=None): + iterators = self._func_execution_context.infer_annotations() + if iterators: + return iterators.iterate(contextualized_node) return self._func_execution_context.get_yield_lazy_values() def py__stop_iteration_returns(self): diff --git a/test/completion/generators.py b/test/completion/generators.py index 566f0036..570ce309 100644 --- a/test/completion/generators.py +++ b/test/completion/generators.py @@ -292,3 +292,22 @@ def test_in_brackets(): x = yield from [1] #? None x + + +# ----------------- +# Annotations +# ----------------- + +from typing import Iterator + +def annotation1() -> float: + yield 1 + +def annotation2() -> Iterator[float]: + yield 1 + + +#? +next(annotation1()) +#? float() +next(annotation2()) diff --git a/test/completion/stdlib.py b/test/completion/stdlib.py index 7dcc6c37..48ab1ee4 100644 --- a/test/completion/stdlib.py +++ b/test/completion/stdlib.py @@ -253,12 +253,34 @@ z.read('name').upper # ----------------- # contextlib # ----------------- - +# python > 2.7 +from typing import Iterator import contextlib with contextlib.closing('asd') as string: #? str() string +@contextlib.contextmanager +def cm1() -> Iterator[float]: + yield 1 +with cm1() as x: + #? float() + x + +@contextlib.contextmanager +def cm2() -> float: + yield 1 +with cm2() as x: + #? + x + +@contextlib.contextmanager +def cm3(): + yield 3 +with cm3() as x: + #? int() + x + # ----------------- # operator # -----------------