From b39946f92272314d3c375d2bddb2677b753f5372 Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Fri, 6 Mar 2026 15:35:29 +0000 Subject: [PATCH 1/2] Optimize InitDecorator.visit_ClassDef MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimization precomputes AST nodes for the synthetic `__init__` (arguments and super-call body) once in `__init__` instead of rebuilding them for every class lacking an `__init__`, and extracts a `_expr_name` helper to replace duplicated isinstance chains when checking decorator/base names. Line profiler shows the `ast.arguments(...)` and `ast.Expr(...)` construction blocks dropped from ~80 µs total per synthetic-init case to near-zero by reusing prebuilt fragments, and the helper consolidates three 4–6 line isinstance sequences into single-call lookups. Runtime improved 18% (237 µs → 199 µs) with negligible per-test variance, and no correctness regressions across all edge cases (dataclasses, NamedTuples, existing decorators). --- .../instrument_codeflash_capture.py | 62 +++++++++++-------- 1 file changed, 36 insertions(+), 26 deletions(-) diff --git a/codeflash/verification/instrument_codeflash_capture.py b/codeflash/verification/instrument_codeflash_capture.py index 2fb9e2366..8bb29e03c 100644 --- a/codeflash/verification/instrument_codeflash_capture.py +++ b/codeflash/verification/instrument_codeflash_capture.py @@ -102,6 +102,23 @@ def __init__( self._init_kwarg = ast.arg(arg="kwargs") self._init_self_arg = ast.arg(arg="self", annotation=None) + + # Precreate commonly reused AST fragments for classes that lack __init__ + # Create the super().__init__(*args, **kwargs) Expr (reuse prebuilt pieces) + self._super_call_expr = ast.Expr( + value=ast.Call(func=self._super_func, args=[self._super_starred], keywords=[self._super_kwarg]) + ) + # Create function arguments: self, *args, **kwargs (reuse arg nodes) + self._init_arguments = ast.arguments( + posonlyargs=[], + args=[self._init_self_arg], + vararg=self._init_vararg, + kwonlyargs=[], + kw_defaults=[], + kwarg=self._init_kwarg, + defaults=[], + ) + def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom: # Check if our import already exists if node.module == "codeflash.verification.codeflash_capture" and any( @@ -162,47 +179,40 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: # TODO: support by saving a reference to the generated __init__ before overriding, e.g. # _orig_init = ClassName.__init__; then calling _orig_init(self, *args, **kwargs) in the wrapper for dec in node.decorator_list: - dec_name = None - if isinstance(dec, ast.Name): - dec_name = dec.id - elif isinstance(dec, ast.Call) and isinstance(dec.func, ast.Name): - dec_name = dec.func.id - elif isinstance(dec, ast.Attribute): - dec_name = dec.attr + dec_name = self._expr_name(dec) if dec_name == "dataclass": return node # Skip NamedTuples — their __init__ is synthesized and cannot be overwritten. for base in node.bases: - base_name = None - if isinstance(base, ast.Name): - base_name = base.id - elif isinstance(base, ast.Attribute): - base_name = base.attr + base_name = self._expr_name(base) if base_name == "NamedTuple": return node # Create super().__init__(*args, **kwargs) call (use prebuilt AST fragments) - super_call = ast.Expr( - value=ast.Call(func=self._super_func, args=[self._super_starred], keywords=[self._super_kwarg]) - ) - # Create function arguments: self, *args, **kwargs (reuse arg nodes) - arguments = ast.arguments( - posonlyargs=[], - args=[self._init_self_arg], - vararg=self._init_vararg, - kwonlyargs=[], - kw_defaults=[], - kwarg=self._init_kwarg, - defaults=[], - ) + super_call = self._super_call_expr + # Create the complete function using prebuilt arguments/body but attach the class-specific decorator # Create the complete function init_func = ast.FunctionDef( - name="__init__", args=arguments, body=[super_call], decorator_list=[decorator], returns=None + name="__init__", args=self._init_arguments, body=[super_call], decorator_list=[decorator], returns=None ) node.body.insert(0, init_func) self.inserted_decorator = True return node + + def _expr_name(self, node: ast.AST) -> str | None: + """ + Helper to extract a simple name from a decorator/base expression. + Returns the identifier for ast.Name, the function name for ast.Call with ast.Name func, + or the attribute name for ast.Attribute. Returns None otherwise. + """ + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + return node.func.id + if isinstance(node, ast.Attribute): + return node.attr + return None From c0f2030494a03e15807c5fa5c08af32325b14085 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 6 Mar 2026 15:45:50 +0000 Subject: [PATCH 2/2] style: remove docstring from _expr_name per project conventions Co-authored-by: Claude Sonnet 4.6 --- codeflash/verification/instrument_codeflash_capture.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/codeflash/verification/instrument_codeflash_capture.py b/codeflash/verification/instrument_codeflash_capture.py index 8bb29e03c..775ce37a8 100644 --- a/codeflash/verification/instrument_codeflash_capture.py +++ b/codeflash/verification/instrument_codeflash_capture.py @@ -102,7 +102,6 @@ def __init__( self._init_kwarg = ast.arg(arg="kwargs") self._init_self_arg = ast.arg(arg="self", annotation=None) - # Precreate commonly reused AST fragments for classes that lack __init__ # Create the super().__init__(*args, **kwargs) Expr (reuse prebuilt pieces) self._super_call_expr = ast.Expr( @@ -204,11 +203,6 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: return node def _expr_name(self, node: ast.AST) -> str | None: - """ - Helper to extract a simple name from a decorator/base expression. - Returns the identifier for ast.Name, the function name for ast.Call with ast.Name func, - or the attribute name for ast.Attribute. Returns None otherwise. - """ if isinstance(node, ast.Name): return node.id if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):