Merge pull request #1867 from codeflash-ai/codeflash/optimize-pr1860-2026-03-18T10.39.52

⚡️ Speed up method `InitDecorator.visit_ClassDef` by 10% in PR #1860 (`fix/attrs-init-instrumentation`)

Author: claude[bot]

codeflash/languages/python/instrument_codeflash_capture.py

@@ -198,6 +198,8 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
                     item.decorator_list.insert(0, decorator)
                     self.inserted_decorator = True
 
+                break
+
         if not has_init:
             # Skip dataclasses — their __init__ is auto-generated at class creation time and isn't in the AST.
             # The synthetic __init__ with super().__init__(*args, **kwargs) overrides it and fails because
@@ -208,25 +210,6 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
                 dec_name = self._expr_name(dec)
                 if dec_name is not None and dec_name.endswith("dataclass"):
                     return node
-
-            # Skip NamedTuples — their __init__ is synthesized and cannot be overwritten.
-            for base in node.bases:
-                base_name = self._expr_name(base)
-                if base_name is not None and base_name.endswith("NamedTuple"):
-                    return node
-
-            # Attrs classes — their __init__ is auto-generated by the decorator at class creation
-            # time. With slots=True (the default for @attrs.define), attrs creates a brand-new class
-            # object, so the __class__ cell baked into a synthesised
-            # `super().__init__(*args, **kwargs)` still refers to the *original* class while `self`
-            # is already an instance of the *new* slots class, causing a TypeError.
-            # We therefore skip modifying the class body and instead emit a module-level
-            # monkey-patch block after the class (handled in visit_Module):
-            #   _codeflash_orig_ClassName_init = ClassName.__init__
-            #   def _codeflash_patched_ClassName_init(self, *a, **kw): ...
-            #   ClassName.__init__ = codeflash_capture(...)(_codeflash_patched_ClassName_init)
-            for dec in node.decorator_list:
-                dec_name = self._expr_name(dec)
                 if dec_name is not None:
                     parts = dec_name.split(".")
                     if len(parts) >= 2 and parts[-2] in _ATTRS_NAMESPACES and parts[-1] in _ATTRS_DECORATOR_NAMES:
@@ -238,6 +221,14 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
                         self.inserted_decorator = True
                         return node
 
+            # Create super().__init__(*args, **kwargs) call (use prebuilt AST fragments)
+
+            # Skip NamedTuples — their __init__ is synthesized and cannot be overwritten.
+            for base in node.bases:
+                base_name = self._expr_name(base)
+                if base_name is not None and base_name.endswith("NamedTuple"):
+                    return node
+
             # Create super().__init__(*args, **kwargs) call (use prebuilt AST fragments)
             super_call = self._super_call_expr
             # Create the complete function using prebuilt arguments/body but attach the class-specific decorator