Fix __future__ imports with leading comments

Author: Otto-AA

mutmut/file_mutation.py

@@ -173,8 +173,8 @@ def combine_mutations_to_source(module: cst.Module, mutations: Sequence[Mutation
     :param mutations: Mutations that should be applied.
     :return: Mutated code and list of mutation names"""
 
-    # add original imports (in particular __future__ imports)
-    result: list[MODULE_STATEMENT] = get_leading_import_statements(module.body)
+    # copy start of the module (in particular __future__ imports)
+    result: list[MODULE_STATEMENT] = get_statements_until_func_or_class(module.body)
     mutation_names: list[str] = []
 
     # statements we still need to potentially mutate and add to the result
@@ -252,17 +252,16 @@ def function_trampoline_arrangement(function: cst.FunctionDef, mutants: Iterable
     return nodes, mutant_names
 
 
-def get_leading_import_statements(statements: Sequence[MODULE_STATEMENT]) -> list[MODULE_STATEMENT]:
-    """Get all `import ...` and `from ... import ...` statements at the start of the module"""
-    leading_import_statements = []
+def get_statements_until_func_or_class(statements: Sequence[MODULE_STATEMENT]) -> list[MODULE_STATEMENT]:
+    """Get all statements until we encounter the first function or class definition"""
+    result = []
 
     for stmt in statements:
-        if m.matches(stmt, m.SimpleStatementLine([m.AtLeastN(matcher=m.Import() | m.ImportFrom(), n=1)])):
-            leading_import_statements.append(stmt)
-        else:
-            break
+        if m.matches(stmt, m.FunctionDef() | m.ClassDef()):
+            return result
+        result.append(stmt)
 
-    return leading_import_statements
+    return result
 
 def group_by_top_level_node(mutations: Sequence[Mutation]) -> Mapping[cst.CSTNode, Sequence[Mutation]]:
     grouped: dict[cst.CSTNode, list[Mutation]] = defaultdict(list)

tests/test_mutation.py

@@ -417,6 +417,20 @@ def foo():
     assert mutated_source.split('\n')[0] == 'from __future__ import annotations'
     assert mutated_source.count('from __future__') == 1
 
+def test_from_future_with_docstring_still_first():
+    source = """
+'''This documents the module'''
+from __future__ import annotations
+from collections.abc import Iterable
+
+def foo():
+    return 1
+""".strip()
+    mutated_source = mutated_module(source)
+    assert mutated_source.split('\n')[0] == "'''This documents the module'''"
+    assert mutated_source.split('\n')[1] == 'from __future__ import annotations'
+    assert mutated_source.count('from __future__') == 1
+
 
 def test_preserve_generators():
     source = '''
@@ -582,10 +596,10 @@ def add(self, value):
 
     assert src == f"""from __future__ import division
 import lib
-{trampoline_impl.strip()}
-{yield_from_trampoline_impl.strip()}
 
 lib.foo()
+{trampoline_impl.strip()}
+{yield_from_trampoline_impl.strip()}
 
 def x_foo__mutmut_orig(a, b):
     return a > b

tests/test_mutmut3.py

@@ -16,9 +16,10 @@ 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')
 
-    expected = trampoline_impl.removesuffix('\n\n') + yield_from_trampoline_impl.removesuffix('\n\n') + """
-a + 1
+    expected = f"""
+a + 1{trampolines}
 
 def x_foo__mutmut_orig(a, b, c):
     return a + b * c
@@ -29,10 +30,10 @@ def x_foo__mutmut_1(a, b, c):
 def x_foo__mutmut_2(a, b, c):
     return a + b / c
 
-x_foo__mutmut_mutants : ClassVar[MutantDict] = {
+x_foo__mutmut_mutants : ClassVar[MutantDict] = {{
 'x_foo__mutmut_1': x_foo__mutmut_1, 
     'x_foo__mutmut_2': x_foo__mutmut_2
-}
+}}
 
 def foo(*args, **kwargs):
     result = _mutmut_trampoline(x_foo__mutmut_orig, x_foo__mutmut_mutants, args, kwargs)