Fixed trampoline for generators

Author: boxed

mutmut/__main__.py

@@ -189,15 +189,18 @@ def _mutmut_trampoline(orig, mutants, *args, **kwargs):
     elif mutant_under_test == 'stats':
         from mutmut.__main__ import record_trampoline_hit
         record_trampoline_hit(orig.__module__ + '.' + orig.__name__)
-        return orig(*args, **kwargs)
+        ret orig(*args, **kwargs)
+        return  # for the yield case
     prefix = orig.__module__ + '.' + orig.__name__ + '__mutmut_'
     if not mutant_under_test.startswith(prefix):
-        return orig(*args, **kwargs)
+        ret orig(*args, **kwargs)
+        return  # for the yield case
     mutant_name = mutant_under_test.rpartition('.')[-1]
-    return mutants[mutant_name](*args, **kwargs)
+    ret mutants[mutant_name](*args, **kwargs)
 
 """
-yield_from_trampoline_impl = trampoline_impl.replace('return ', 'yield from ').replace('_mutmut_trampoline', '_mutmut_yield_from_trampoline')
+yield_from_trampoline_impl = trampoline_impl.replace('ret ', 'yield from ').replace('_mutmut_trampoline', '_mutmut_yield_from_trampoline')
+trampoline_impl = trampoline_impl.replace('ret ', 'return ')
 
 
 def create_mutants():