Make generators return more correct values with while loops, fixes #683

jedi/inference/lazy_value.py

@@ -3,8 +3,10 @@
 
 
 class AbstractLazyValue(object):
-    def __init__(self, data):
+    def __init__(self, data, min=1, max=1):
         self.data = data
+        self.min = min
+        self.max = max
 
     def __repr__(self):
         return '<%s: %s>' % (self.__class__.__name__, self.data)
@@ -26,16 +28,16 @@ def infer(self):
 
 
 class LazyUnknownValue(AbstractLazyValue):
-    def __init__(self):
-        super(LazyUnknownValue, self).__init__(None)
+    def __init__(self, min=1, max=1):
+        super(LazyUnknownValue, self).__init__(None, min, max)
 
     def infer(self):
         return NO_VALUES
 
 
 class LazyTreeValue(AbstractLazyValue):
-    def __init__(self, context, node):
-        super(LazyTreeValue, self).__init__(node)
+    def __init__(self, context, node, min=1, max=1):
+        super(LazyTreeValue, self).__init__(node, min, max)
         self.context = context
         # We need to save the predefined names. It's an unfortunate side effect
         # that needs to be tracked otherwise results will be wrong.

jedi/inference/syntax_tree.py

@@ -797,7 +797,8 @@ def check_tuple_assignments(name, value_set):
         if isinstance(index, slice):
             # For no star unpacking is not possible.
             return NO_VALUES
-        for _ in range(index + 1):
+        i = 0
+        while i <= index:
             try:
                 lazy_value = next(iterated)
             except StopIteration:
@@ -806,6 +807,8 @@ def check_tuple_assignments(name, value_set):
                 # index number is high. Therefore break if the loop is
                 # finished.
                 return NO_VALUES
+            else:
+                i += lazy_value.max
         value_set = lazy_value.infer()
     return value_set
 

jedi/inference/value/function.py

@@ -265,7 +265,7 @@ def get_yield_lazy_values(self, is_async=False):
             else:
                 types = self.get_return_values(check_yields=True)
                 if types:
-                    yield LazyKnownValues(types)
+                    yield LazyKnownValues(types, min=0, max=float('inf'))
                 return
             last_for_stmt = for_stmt
 

test/completion/generators.py

@@ -156,7 +156,7 @@ def simple():
 #? int() str()
 a
 # For now this is ok.
-#?
+#? int() str()
 b