Fix enums with custom __init__ failing to import under mutmut

src/mutmut/mutation/file_mutation.py

@@ -35,6 +35,8 @@
 from mutmut.utils.format_utils import is_mutated_method_name
 
 NEVER_MUTATE_FUNCTION_NAMES = {"__getattribute__", "__setattr__", "__new__"}
+NEVER_MUTATE_ENUM_FUNCTION_NAMES = {"__init__", "_generate_next_value_"}
+ENUM_BASE_NAMES = {"Enum", "IntEnum", "StrEnum", "Flag", "IntFlag"}
 NEVER_MUTATE_FUNCTION_CALLS = {"len", "isinstance"}
 
 
@@ -217,6 +219,7 @@ def __init__(
         )
         self.ignored_classes: set[str] = set()
         self.ignored_functions: set[str] = set()
+        self._enum_creation_hooks: set[cst.FunctionDef] = set()
 
     def on_visit(self, node: cst.CSTNode) -> bool:
         if self._skip_node_and_children(node):
@@ -300,6 +303,15 @@ def _skip_node_and_children(self, node: cst.CSTNode) -> bool:
             if isinstance(node, cst.BaseExpression):
                 return True
 
+        if isinstance(node, cst.ClassDef) and _is_enum_class(node):
+            self._enum_creation_hooks.update(
+                child
+                for child in node.body.body
+                if isinstance(child, cst.FunctionDef) and child.name.value in NEVER_MUTATE_ENUM_FUNCTION_NAMES
+            )
+        if node in self._enum_creation_hooks:
+            return True
+
         if (
             isinstance(node, cst.Call)
             and isinstance(node.func, cst.Name)
@@ -344,6 +356,11 @@ def _skip_node_and_children(self, node: cst.CSTNode) -> bool:
         return False
 
 
+def _is_enum_class(node: cst.ClassDef) -> bool:
+    bases = (base.value.attr if isinstance(base.value, cst.Attribute) else base.value for base in node.bases)
+    return any(isinstance(base, cst.Name) and base.value in ENUM_BASE_NAMES for base in bases)
+
+
 def _walk(node: cst.CSTNode) -> Iterator[cst.CSTNode]:
     """Yield `node` and every node below it."""
     stack = [node]

tests/mutation/test_mutation_runtime.py

@@ -38,6 +38,39 @@ def describe(self):
     monkeypatch.setenv("MUTANT_UNDER_TEST", mutant_name)
 
 
+def test_enum_creation_hooks_runtime_execution(monkeypatch):
+    """Test that an enum whose __init__ and _generate_next_value_ run during class creation can be executed."""
+    source = """
+from enum import Enum, auto
+
+class Color(Enum):
+    def _generate_next_value_(name, start, count, last_values):
+        return count + 10
+
+    RED = auto()
+    GREEN = auto()
+
+    def __init__(self, value):
+        self.label = f"c{value}"
+
+    def describe(self):
+        return self.name.lower()
+""".strip()
+
+    mutated_file = mutate_file_contents("test.py", source)
+    mutated_code, mutant_names = mutated_file.code, mutated_file.mutant_names
+    assert len(mutant_names) > 0, "Should have at least one mutant"
+
+    monkeypatch.setenv("MUTANT_UNDER_TEST", "none")
+    namespace = {"__name__": "test_module"}
+    exec(mutated_code, namespace)
+    Color = namespace["Color"]
+
+    assert Color.GREEN.value == 11
+    assert Color.GREEN.label == "c11"
+    assert Color.GREEN.describe() == "green"
+
+
 def test_type_annotation_runtime_execution(monkeypatch):
     """Test that forward and quoted annotations can be executed