Support async generators (#407)

Functions with "yield" return a generator
which we can simply return as a value to
the caller.

Author: Otto-AA

e2e_projects/my_lib/pyproject.toml

@@ -17,4 +17,7 @@ dev = [
 ]
 
 [tool.mutmut]
-debug = true
+debug = true
+
+[tool.pytest]
+asyncio_default_fixture_loop_scope = "function"

e2e_projects/my_lib/src/my_lib/__init__.py

@@ -26,6 +26,29 @@ def fibonacci(n: int) -> int:
         return n
     return fibonacci(n - 1) + fibonacci(n - 2)
 
+async def async_consumer():
+    results = []
+    async for i in async_generator():
+        results.append(i)
+    return results
+
+async def async_generator():
+    for i in range(10):
+        yield i
+
+def simple_consumer():
+    generator = double_generator()
+    next(generator) # skip the initial yield
+    results = []
+    for i in range(10):
+        results.append(generator.send(i))
+    return results
+
+def double_generator():
+    while True:
+        x = yield
+        yield x * 2
+
 @cache
 def cached_fibonacci(n: int) -> int:
     if n <= 1:

e2e_projects/my_lib/tests/test_my_lib.py

@@ -1,4 +1,5 @@
-from my_lib import hello, Point, badly_tested, make_greeter, fibonacci, cached_fibonacci, escape_sequences
+from my_lib import hello, Point, badly_tested, make_greeter, fibonacci, cached_fibonacci, escape_sequences, simple_consumer, async_consumer
+import pytest
 
 """These tests are flawed on purpose, some mutants survive and some are killed."""
 
@@ -33,4 +34,14 @@ def test_fibonacci():
     assert cached_fibonacci(1) == 1
 
 def test_escape_sequences():
-    assert escape_sequences().lower() == "foofoo\\\'\"\a\b\f\n\r\t\v\111\x10\N{ghost}\u1234\U0001F51F".lower()
+    assert escape_sequences().lower() == "foofoo\\\'\"\a\b\f\n\r\t\v\111\x10\N{ghost}\u1234\U0001F51F".lower()
+
+def test_simple_consumer():
+    # only verifying length, should report surviving mutants for the contents
+    assert len(simple_consumer()) == 10
+
+@pytest.mark.asyncio
+async def test_async_consumer():
+    result = await async_consumer()
+    assert result == list(range(10))
+

mutmut/file_mutation.py

@@ -8,7 +8,7 @@
 import libcst as cst
 from libcst.metadata import PositionProvider, MetadataWrapper
 import libcst.matchers as m
-from mutmut.trampoline_templates import build_trampoline, mangle_function_name, trampoline_impl, yield_from_trampoline_impl
+from mutmut.trampoline_templates import build_trampoline, mangle_function_name, trampoline_impl
 from mutmut.node_mutation import mutation_operators, OPERATORS_TYPE
 
 NEVER_MUTATE_FUNCTION_NAMES = { "__getattribute__", "__setattr__", "__new__" }
@@ -165,8 +165,6 @@ def _skip_node_and_children(self, node: cst.CSTNode):
 # convert str trampoline implementations to CST nodes with some whitespace
 trampoline_impl_cst = list(cst.parse_module(trampoline_impl).body)
 trampoline_impl_cst[-1] = trampoline_impl_cst[-1].with_changes(leading_lines = [cst.EmptyLine(), cst.EmptyLine()])
-yield_from_trampoline_impl_cst = list(cst.parse_module(yield_from_trampoline_impl).body)
-yield_from_trampoline_impl_cst[-1] = yield_from_trampoline_impl_cst[-1].with_changes(leading_lines = [cst.EmptyLine(), cst.EmptyLine()])
 
 
 def combine_mutations_to_source(module: cst.Module, mutations: Sequence[Mutation]) -> tuple[str, Sequence[str]]:
@@ -185,7 +183,6 @@ def combine_mutations_to_source(module: cst.Module, mutations: Sequence[Mutation
 
     # trampoline functions
     result.extend(trampoline_impl_cst)
-    result.extend(yield_from_trampoline_impl_cst)
 
     mutations_within_function = group_by_top_level_node(mutations)
 
@@ -234,7 +231,6 @@ def function_trampoline_arrangement(function: cst.FunctionDef, mutants: Iterable
 
     name = function.name.value
     mangled_name = mangle_function_name(name=name, class_name=class_name) + '__mutmut'
-    _is_generator = is_generator(function)
 
     # copy of original function
     nodes.append(function.with_changes(name=cst.Name(mangled_name + '_orig')))
@@ -248,7 +244,7 @@ def function_trampoline_arrangement(function: cst.FunctionDef, mutants: Iterable
         nodes.append(mutated_method) # type: ignore
 
     # trampoline that forwards the calls
-    trampoline = list(cst.parse_module(build_trampoline(orig_name=name, mutants=mutant_names, class_name=class_name, is_generator=_is_generator)).body)
+    trampoline = list(cst.parse_module(build_trampoline(orig_name=name, mutants=mutant_names, class_name=class_name)).body)
     trampoline[0] = trampoline[0].with_changes(leading_lines=[cst.EmptyLine()])
     nodes.extend(trampoline)
 
@@ -274,28 +270,6 @@ def group_by_top_level_node(mutations: Sequence[Mutation]) -> Mapping[cst.CSTNod
 
     return grouped
 
-def is_generator(function: cst.FunctionDef) -> bool:
-    """Return True if the function has yield statement(s)."""
-    visitor = IsGeneratorVisitor(function)
-    function.visit(visitor)
-    return visitor.is_generator
-
-class IsGeneratorVisitor(cst.CSTVisitor):
-    """Check if a function is a generator.
-    We do so by checking if any child is a Yield statement, but not looking into inner function definitions."""
-    def __init__(self, original_function: cst.FunctionDef):
-        self.is_generator = False
-        self.original_function: cst.FunctionDef = original_function
-
-    def visit_FunctionDef(self, node):
-        # do not recurse into inner function definitions
-        if self.original_function != node:
-            return False
-
-    def visit_Yield(self, node):
-        self.is_generator = True
-        return False
-
 def pragma_no_mutate_lines(source: str) -> set[int]:
     return {
         i + 1

mutmut/trampoline_templates.py

@@ -1,6 +1,6 @@
 CLASS_NAME_SEPARATOR = 'ǁ'
 
-def build_trampoline(*, orig_name, mutants, class_name, is_generator):
+def build_trampoline(*, orig_name, mutants, class_name):
     mangled_name = mangle_function_name(name=orig_name, class_name=class_name)
 
     mutants_dict = f'{mangled_name}__mutmut_mutants : ClassVar[MutantDict] = {{\n' + ', \n    '.join(f'{repr(m)}: {m}' for m in mutants) + '\n}'
@@ -12,18 +12,13 @@ def build_trampoline(*, orig_name, mutants, class_name, is_generator):
         access_suffix = '")'
         self_arg = ', self'
 
-    if is_generator:
-        yield_statement = 'yield from '  # note the space at the end!
-        trampoline_name = '_mutmut_yield_from_trampoline'
-    else:
-        yield_statement = ''
-        trampoline_name = '_mutmut_trampoline'
+    trampoline_name = '_mutmut_trampoline'
 
     return f"""
 {mutants_dict}
 
 def {orig_name}({'self, ' if class_name is not None else ''}*args, **kwargs):
-    result = {yield_statement}{trampoline_name}({access_prefix}{mangled_name}__mutmut_orig{access_suffix}, {access_prefix}{mangled_name}__mutmut_mutants{access_suffix}, args, kwargs{self_arg})
+    result = {trampoline_name}({access_prefix}{mangled_name}__mutmut_orig{access_suffix}, {access_prefix}{mangled_name}__mutmut_mutants{access_suffix}, args, kwargs{self_arg})
     return result 
 
 {orig_name}.__signature__ = _mutmut_signature({mangled_name}__mutmut_orig)
@@ -62,11 +57,11 @@ def _mutmut_trampoline(orig, mutants, call_args, call_kwargs, self_arg = None):
         from mutmut.__main__ import record_trampoline_hit
         record_trampoline_hit(orig.__module__ + '.' + orig.__name__)
         result = orig(*call_args, **call_kwargs)
-        return result  # for the yield case
+        return result
     prefix = orig.__module__ + '.' + orig.__name__ + '__mutmut_'
     if not mutant_under_test.startswith(prefix):
         result = orig(*call_args, **call_kwargs)
-        return result  # for the yield case
+        return result
     mutant_name = mutant_under_test.rpartition('.')[-1]
     if self_arg:
         # call to a class method where self is not bound
@@ -76,4 +71,3 @@ def _mutmut_trampoline(orig, mutants, call_args, call_kwargs, self_arg = None):
     return result
 
 """
-yield_from_trampoline_impl = trampoline_impl.replace('result = ', 'result = yield from ').replace('_mutmut_trampoline', '_mutmut_yield_from_trampoline')

test_requirements.txt

@@ -1,4 +1,5 @@
 pytest
+pytest-asyncio>=1.0.0
 mock>=2.0.0
 coverage
 whatthepatch==0.0.6

tests/e2e/snapshots/my_lib.json

@@ -25,6 +25,21 @@
     "my_lib.x_fibonacci__mutmut_7": 0,
     "my_lib.x_fibonacci__mutmut_8": 0,
     "my_lib.x_fibonacci__mutmut_9": 0,
+    "my_lib.x_async_consumer__mutmut_1": 1,
+    "my_lib.x_async_consumer__mutmut_2": 1,
+    "my_lib.x_async_generator__mutmut_1": 1,
+    "my_lib.x_async_generator__mutmut_2": 1,
+    "my_lib.x_simple_consumer__mutmut_1": 1,
+    "my_lib.x_simple_consumer__mutmut_2": 1,
+    "my_lib.x_simple_consumer__mutmut_3": 1,
+    "my_lib.x_simple_consumer__mutmut_4": 1,
+    "my_lib.x_simple_consumer__mutmut_5": 1,
+    "my_lib.x_simple_consumer__mutmut_6": 0,
+    "my_lib.x_simple_consumer__mutmut_7": 1,
+    "my_lib.x_double_generator__mutmut_1": 1,
+    "my_lib.x_double_generator__mutmut_2": 1,
+    "my_lib.x_double_generator__mutmut_3": 0,
+    "my_lib.x_double_generator__mutmut_4": 0,
     "my_lib.x\u01c1Point\u01c1__init____mutmut_1": 1,
     "my_lib.x\u01c1Point\u01c1__init____mutmut_2": 1,
     "my_lib.x\u01c1Point\u01c1abs__mutmut_1": 33,

tests/test_mutation.py

@@ -13,8 +13,8 @@
     MutmutProgrammaticFailException,
     CatchOutput,
 )
-from mutmut.trampoline_templates import trampoline_impl, yield_from_trampoline_impl, mangle_function_name
-from mutmut.file_mutation import create_mutations, mutate_file_contents, is_generator
+from mutmut.trampoline_templates import trampoline_impl, mangle_function_name
+from mutmut.file_mutation import create_mutations, mutate_file_contents
 
 def mutants_for_source(source: str) -> list[str]:
     module, mutated_nodes = create_mutations(source)
@@ -514,43 +514,6 @@ def foo():
     assert mutated_source.count('from __future__') == 1
 
 
-def test_preserve_generators():
-    source = '''
-    def foo():
-        yield 1
-    '''.strip()
-    mutated_source = mutated_module(source)
-    assert 'yield from _mutmut_yield_from_trampoline' in mutated_source
-
-
-def test_is_generator():
-    source = '''
-    def foo():
-        yield 1
-    '''.strip()
-    assert is_generator(parse_statement(source)) # type: ignore
-
-    source = '''
-    def foo():
-        yield from bar()
-    '''.strip()
-    assert is_generator(parse_statement(source)) # type: ignore
-
-    source = '''
-    def foo():
-        return 1
-    '''.strip()
-    assert not is_generator(parse_statement(source)) # type: ignore
-
-    source = '''
-    def foo():
-        def bar():
-            yield 2
-        return 1
-    '''.strip()
-    assert not is_generator(parse_statement(source)) # type: ignore
-
-
 # Negate the effects of CatchOutput because it does not play nicely with capfd in GitHub Actions
 @patch.object(CatchOutput, 'dump_output')
 @patch.object(CatchOutput, 'stop')
@@ -678,7 +641,6 @@ def add(self, value):
 
 lib.foo()
 {trampoline_impl.strip()}
-{yield_from_trampoline_impl.strip()}
 
 def x_foo__mutmut_orig(a, b):
     return a > b
@@ -708,7 +670,7 @@ def x_bar__mutmut_1():
 }}
 
 def bar(*args, **kwargs):
-    result = yield from _mutmut_yield_from_trampoline(x_bar__mutmut_orig, x_bar__mutmut_mutants, args, kwargs)
+    result = _mutmut_trampoline(x_bar__mutmut_orig, x_bar__mutmut_mutants, args, kwargs)
     return result 
 
 bar.__signature__ = _mutmut_signature(x_bar__mutmut_orig)

tests/test_mutmut3.py

@@ -1,7 +1,4 @@
-from mutmut.trampoline_templates import (
-    trampoline_impl,
-    yield_from_trampoline_impl,
-)
+from mutmut.trampoline_templates import trampoline_impl
 from mutmut.file_mutation import mutate_file_contents
 
 def mutated_module(source: str) -> str:
@@ -16,7 +13,7 @@ def test_mutate_file_contents():
 def foo(a, b, c):
     return a + b * c
 """
-    trampolines = trampoline_impl.removesuffix('\n\n') + yield_from_trampoline_impl.removesuffix('\n\n')
+    trampolines = trampoline_impl.removesuffix('\n\n')
 
     expected = f"""
 a + 1{trampolines}
@@ -54,7 +51,7 @@ def foo(a: List[int]) -> int:
     return 1
 """
 
-    expected = trampoline_impl.removesuffix('\n\n') + yield_from_trampoline_impl.removesuffix('\n\n') + """
+    expected = trampoline_impl.removesuffix('\n\n') + """
 def x_foo__mutmut_orig(a: List[int]) -> int:
     return 1
 def x_foo__mutmut_1(a: List[int]) -> int: