From e930f478611d86e29ed0353b0421ee1c861d81f1 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 29 Jan 2020 10:13:46 +0100
Subject: [PATCH] Make generators return more correct values with while loops,
 fixes #683

---
 jedi/inference/lazy_value.py     | 12 +++++++-----
 jedi/inference/syntax_tree.py    |  5 ++++-
 jedi/inference/value/function.py |  2 +-
 test/completion/generators.py    |  2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/jedi/inference/lazy_value.py b/jedi/inference/lazy_value.py
index fa2f96096..44f8d5967 100644
--- a/jedi/inference/lazy_value.py
+++ b/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.
diff --git a/jedi/inference/syntax_tree.py b/jedi/inference/syntax_tree.py
index 6512fa76b..6411fc818 100644
--- a/jedi/inference/syntax_tree.py
+++ b/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
 
diff --git a/jedi/inference/value/function.py b/jedi/inference/value/function.py
index ab9e65c75..ae28ae9aa 100644
--- a/jedi/inference/value/function.py
+++ b/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
 
diff --git a/test/completion/generators.py b/test/completion/generators.py
index 5ea3082e5..e36f39287 100644
--- a/test/completion/generators.py
+++ b/test/completion/generators.py
@@ -156,7 +156,7 @@ def simple():
 #? int() str()
 a
 # For now this is ok.
-#?
+#? int() str()
 b