From a06ca5d0354f9447ee8b004301cca80b4a18b94d Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Mon, 24 Sep 2018 20:59:43 +0200
Subject: [PATCH] Fix generator return issues

---
 jedi/evaluate/context/iterable.py |  7 +++++++
 jedi/evaluate/context/typing.py   | 20 +++++++++++---------
 jedi/evaluate/syntax_tree.py      |  7 ++++---
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/jedi/evaluate/context/iterable.py b/jedi/evaluate/context/iterable.py
index 0d7c9e44..0c4b0137 100644
--- a/jedi/evaluate/context/iterable.py
+++ b/jedi/evaluate/context/iterable.py
@@ -57,12 +57,19 @@ class GeneratorBase(BuiltinOverwrite, IterableMixin):
             .execute_annotation()
         return generator
 
+    @publish_method('__iter__')
+    def py__iter__(self):
+        return ContextSet([self])
+
     @publish_method('send')
     @publish_method('next', python_version_match=2)
     @publish_method('__next__', python_version_match=3)
     def py__next__(self):
         return ContextSet.from_sets(lazy_context.infer() for lazy_context in self.py__iter__())
 
+    def py__stop_iteration_returns(self):
+        return ContextSet([compiled.builtin_from_name(self.evaluator, u'None')])
+
     @property
     def name(self):
         return compiled.CompiledContextName(self, 'generator')
diff --git a/jedi/evaluate/context/typing.py b/jedi/evaluate/context/typing.py
index a71840ad..2ac86ca7 100644
--- a/jedi/evaluate/context/typing.py
+++ b/jedi/evaluate/context/typing.py
@@ -6,7 +6,7 @@ contexts.
 from jedi._compatibility import unicode
 from jedi import debug
 from jedi.evaluate.cache import evaluator_method_cache
-from jedi.evaluate.compiled import builtin_from_name, CompiledObject
+from jedi.evaluate.compiled import builtin_from_name
 from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS, Context, \
     iterator_to_context_set, HelperContextMixin, ContextWrapper
 from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext
@@ -17,7 +17,7 @@ from jedi.evaluate.filters import FilterWrapper, NameWrapper, \
     AbstractTreeName, AbstractNameDefinition, ContextName
 from jedi.evaluate.helpers import is_string
 from jedi.evaluate.imports import Importer
-from jedi.evaluate.context import ClassContext
+from jedi.evaluate.context.klass import py__mro__, ClassContext
 
 _PROXY_CLASS_TYPES = 'Tuple Generic Protocol Callable Type'.split()
 _TYPE_ALIAS_TYPES = {
@@ -607,11 +607,13 @@ class LazyAnnotatedBaseClass(object):
 
 class InstanceWrapper(ContextWrapper):
     def py__stop_iteration_returns(self):
-        cls = self._wrapped_context.class_context
-        if cls.py__name__() == 'Generator':
-            given_types = cls.get_given_types()
-            try:
-                return given_types[2].execute_annotation()
-            except IndexError:
-                pass
+        for cls in py__mro__(self._wrapped_context.class_context):
+            if cls.py__name__() == 'Generator':
+                given_types = cls.get_given_types()
+                try:
+                    return given_types[2].execute_annotation()
+                except IndexError:
+                    pass
+            elif cls.py__name__() == 'Iterator':
+                return ContextSet([builtin_from_name(self.evaluator, u'None')])
         return self._wrapped_context.py__stop_iteration_returns()
diff --git a/jedi/evaluate/syntax_tree.py b/jedi/evaluate/syntax_tree.py
index 6d06f277..19b36b57 100644
--- a/jedi/evaluate/syntax_tree.py
+++ b/jedi/evaluate/syntax_tree.py
@@ -100,7 +100,7 @@ def eval_node(context, element):
             await_context_set = context_set.py__getattribute__(u"__await__")
             if not await_context_set:
                 debug.warning('Tried to run py__await__ on context %s', context)
-            return _py__stop_iteration_returns(await_context_set.execute_evaluated())
+            return await_context_set.execute_evaluated().py__stop_iteration_returns()
         return context_set
     elif typ in ('testlist_star_expr', 'testlist',):
         # The implicit tuple in statements.
@@ -136,8 +136,9 @@ def eval_node(context, element):
         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)
-            return _py__stop_iteration_returns(generators)
+            generators = context.eval_node(element) \
+                .py__getattribute__('__iter__').execute_evaluated()
+            return generators.py__stop_iteration_returns()
 
         # Generator.send() is not implemented.
         return NO_CONTEXTS