From 0d0213ee4caba5020ae69e94f56b60aac4f42bb4 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 28 Feb 2018 22:35:58 +0100
Subject: [PATCH] Support generator returns when used with yield from.

---
 jedi/evaluate/context/iterable.py | 14 +++++++++----
 jedi/evaluate/syntax_tree.py      | 17 ++++++++++++++++
 test/completion/generators.py     | 33 +++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/jedi/evaluate/context/iterable.py b/jedi/evaluate/context/iterable.py
index 9de6266f..7a17bc75 100644
--- a/jedi/evaluate/context/iterable.py
+++ b/jedi/evaluate/context/iterable.py
@@ -42,12 +42,11 @@ from jedi.parser_utils import get_comp_fors
 
 
 class AbstractIterableMixin(object):
-    @property
-    def name(self):
-        return compiled.CompiledContextName(self, self.array_type)
+    def py__stop_iteration_returns(self):
+        return ContextSet(compiled.builtin_from_name(self.evaluator, u'None'))
 
 
-class GeneratorBase(BuiltinOverwrite):
+class GeneratorBase(BuiltinOverwrite, AbstractIterableMixin):
     array_type = None
     special_object_identifier = u'GENERATOR_OBJECT'
 
@@ -71,6 +70,9 @@ class Generator(GeneratorBase):
     def py__iter__(self):
         return self._func_execution_context.get_yield_lazy_contexts()
 
+    def py__stop_iteration_returns(self):
+        return self._func_execution_context.get_return_values()
+
     def __repr__(self):
         return "<%s of %s>" % (type(self).__name__, self._func_execution_context)
 
@@ -182,6 +184,10 @@ class ComprehensionMixin(object):
 class Sequence(BuiltinOverwrite, AbstractIterableMixin):
     api_type = u'instance'
 
+    @property
+    def name(self):
+        return compiled.CompiledContextName(self, self.array_type)
+
     @memoize_method
     def get_object(self):
         compiled_obj = compiled.builtin_from_name(self.evaluator, self.array_type)
diff --git a/jedi/evaluate/syntax_tree.py b/jedi/evaluate/syntax_tree.py
index 7a131647..2ea81409 100644
--- a/jedi/evaluate/syntax_tree.py
+++ b/jedi/evaluate/syntax_tree.py
@@ -124,6 +124,23 @@ def eval_node(context, element):
         return eval_node(context, element.children[0])
     elif typ == 'annassign':
         return pep0484._evaluate_for_annotation(context, element.children[1])
+    elif typ == 'yield_expr':
+        if len(element.children) and element.children[1].type == 'yield_arg':
+            # Implies that it's a yield from.
+            element = element.children[1].children[1]
+            generators = context.eval_node(element)
+            results = ContextSet()
+            for generator in generators:
+                try:
+                    method = generator.py__stop_iteration_returns
+                except AttributeError:
+                    debug.warning('%s is not actually a generator', generator)
+                else:
+                    results |= method()
+            return results
+
+        # Generator.send() is not implemented.
+        return NO_CONTEXTS
     else:
         return eval_or_test(context, element)
 
diff --git a/test/completion/generators.py b/test/completion/generators.py
index 2327b185..3e44b4c6 100644
--- a/test/completion/generators.py
+++ b/test/completion/generators.py
@@ -223,9 +223,42 @@ next(yield_from())
 def yield_from_multiple():
     yield from iter([1])
     yield str()
+    return 2.0
 
 x, y = yield_from_multiple()
 #? int()
 x
 #? str()
 y
+
+def test_nested():
+    x = yield from yield_from_multiple()
+    #? float()
+    x
+    yield x
+
+x, y, z = yield_nested()
+#? int()
+x
+#? str()
+y
+# For whatever reason this is currently empty
+#?
+z
+
+
+def test_in_brackets():
+    x = 1 + (yield from yield_from_multiple())
+    #? float()
+    x
+
+    generator = (1 for 1 in [1])
+    x = yield from generator
+    #? None
+    x
+    x = yield from 1
+    #?
+    x
+    x = yield from [1]
+    #? None
+    x