From d3c437e891be4f59735add0515e5e479243ffa0b Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sat, 7 Jan 2017 12:43:15 +0100
Subject: [PATCH] Restructure yield code to make it less error prone.

---
 jedi/evaluate/recursion.py      | 28 ++++++++++++++++------------
 jedi/evaluate/representation.py | 11 ++++++-----
 sith.py                         |  2 +-
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/jedi/evaluate/recursion.py b/jedi/evaluate/recursion.py
index ed05e60f..088fc93c 100644
--- a/jedi/evaluate/recursion.py
+++ b/jedi/evaluate/recursion.py
@@ -27,7 +27,7 @@ def execution_allowed(evaluator, node):
     pushed_nodes = evaluator.recursion_detector.pushed_nodes
 
     if node in pushed_nodes:
-        debug.warning('catched stmt recursion: %s against %s @%s', node,
+        debug.warning('catched stmt recursion: %s @%s', node,
                       node.start_pos)
         yield False
     else:
@@ -36,17 +36,21 @@ def execution_allowed(evaluator, node):
         pushed_nodes.pop()
 
 
-def execution_recursion_decorator(func):
-    def run(execution, **kwargs):
-        detector = execution.evaluator.execution_recursion_detector
-        if detector.push_execution(execution):
-            result = set()
-        else:
-            result = func(execution, **kwargs)
-        detector.pop_execution()
-        return result
-
-    return run
+def execution_recursion_decorator(default=set()):
+    def decorator(func):
+        def wrapper(execution, **kwargs):
+            detector = execution.evaluator.execution_recursion_detector
+            allowed = detector.push_execution(execution)
+            try:
+                if allowed:
+                    result = default
+                else:
+                    result = func(execution, **kwargs)
+            finally:
+                detector.pop_execution()
+            return result
+        return wrapper
+    return decorator
 
 
 class ExecutionRecursionDetector(object):
diff --git a/jedi/evaluate/representation.py b/jedi/evaluate/representation.py
index 7f5aa187..e4107501 100644
--- a/jedi/evaluate/representation.py
+++ b/jedi/evaluate/representation.py
@@ -295,7 +295,7 @@ class FunctionExecutionContext(Executed):
         self.tree_node = function_context.tree_node
 
     @memoize_default(default=set())
-    @recursion.execution_recursion_decorator
+    @recursion.execution_recursion_decorator()
     def get_return_values(self, check_yields=False):
         funcdef = self.tree_node
         if funcdef.type == 'lambda':
@@ -321,8 +321,6 @@ class FunctionExecutionContext(Executed):
             if check is flow_analysis.REACHABLE:
                 debug.dbg('Return reachable: %s', r)
                 break
-        if check_yields:
-            return context.get_merged_lazy_context(list(types))
         return types
 
     def _eval_yield(self, yield_expr):
@@ -334,7 +332,7 @@ class FunctionExecutionContext(Executed):
         else:
             yield context.LazyTreeContext(self, node)
 
-    @recursion.execution_recursion_decorator
+    @recursion.execution_recursion_decorator(default=iter([]))
     def get_yield_values(self):
         for_parents = [(y, tree.search_ancestor(y, ('for_stmt', 'funcdef',
                                                     'while_stmt', 'if_stmt')))
@@ -358,7 +356,9 @@ class FunctionExecutionContext(Executed):
             elif for_stmt == self.tree_node:
                 yields_order.append((None, [yield_]))
             else:
-                yield self.get_return_values(check_yields=True)
+                types = self.get_return_values(check_yields=True)
+                if types:
+                    yield context.get_merged_lazy_context(list(types))
                 return
             last_for_stmt = for_stmt
 
@@ -373,6 +373,7 @@ class FunctionExecutionContext(Executed):
                 input_node = for_stmt.get_input_node()
                 for_types = self.eval_node(input_node)
                 ordered = iterable.py__iter__(evaluator, for_types, input_node)
+                ordered = list(ordered)
                 for lazy_context in ordered:
                     dct = {str(for_stmt.children[1]): lazy_context.infer()}
                     with helpers.predefine_names(self, for_stmt, dct):
diff --git a/sith.py b/sith.py
index 903ef0ef..2f676220 100755
--- a/sith.py
+++ b/sith.py
@@ -173,7 +173,7 @@ class TestCase(object):
                 self.show_location(completion.line, completion.column)
 
     def show_errors(self):
-        print(self.traceback)
+        sys.stderr.write(self.traceback)
         print(("Error with running Script(...).{operation}() with\n"
               "\tpath:   {path}\n"
               "\tline:   {line}\n"