Fix issues with generators, fixes #1624

jedi/inference/value/function.py

@@ -202,7 +202,7 @@ def name(self):
 
 
 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 @@ def get_return_values(self, check_yields=False):
             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 @@ def get_filters(self, until_position=None, origin_scope=None):
             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 @@ def get_param_names(self):
 
 
 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

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 @@ def py__stop_iteration_returns(self):
     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 @@ def __init__(self, inference_state, func_execution_context):
         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):

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())

test/completion/stdlib.py

@@ -253,12 +253,34 @@ def shout(self): pass
 # -----------------
 # 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
 # -----------------