Merge branch 'main' into multi-language

.github/workflows/mypy.yml

@@ -22,11 +22,10 @@ jobs:
 
       - name: Install uv
         uses: astral-sh/setup-uv@v6
-        with:
-          version: "0.5.30"
 
       - name: sync uv
         run: |
+          uv venv --seed
           uv sync
 
 

.gitignore

@@ -262,3 +262,6 @@ tessl.json
 **/node_modules/**
 /dist-nuitka/main.dist/*
 packages/codeflash/.npmrc
+
+# Tessl auto-generates AGENTS.md on install; ignore to avoid cluttering git status
+AGENTS.md

CLAUDE.md

@@ -88,6 +88,10 @@ else:
 - Commit message body should be concise (1-2 sentences max)
 - PR titles should also use conventional format
 
+<!-- Section below is auto-generated by `tessl install` - do not edit manually -->
+
 # Agent Rules <!-- tessl-managed -->
 
 @.tessl/RULES.md follow the [instructions](.tessl/RULES.md)
+
+@AGENTS.md

codeflash/LICENSE

@@ -3,7 +3,7 @@ Business Source License 1.1
 Parameters
 
 Licensor:             CodeFlash Inc.
-Licensed Work:        Codeflash Client version 0.19.x
+Licensed Work:        Codeflash Client version 0.20.x
                       The Licensed Work is (c) 2024 CodeFlash Inc.
 
 Additional Use Grant: None. Production use of the Licensed Work is only permitted
@@ -13,7 +13,7 @@ Additional Use Grant: None. Production use of the Licensed Work is only permitte
                       Platform. Please visit codeflash.ai for further
                       information.
 
-Change Date:          2029-12-21
+Change Date:          2030-01-26
 
 Change License:       MIT
 

codeflash/api/aiservice.py

@@ -268,7 +268,7 @@ class AiServiceClient:
         logger.info("!lsp|Rewriting as a JIT function…")
         console.rule()
         try:
-            response = self.make_ai_service_request("/rewrite_jit", payload=payload, timeout=60)
+            response = self.make_ai_service_request("/rewrite_jit", payload=payload, timeout=self.timeout)
         except requests.exceptions.RequestException as e:
             logger.exception(f"Error generating jit rewritten candidate: {e}")
             ph("cli-jit-rewrite-error-caught", {"error": str(e)})

codeflash/cli_cmds/cli.py

@@ -84,6 +84,9 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--no-gen-tests", action="store_true", help="Do not generate tests, use only existing tests for optimization."
     )
+    parser.add_argument(
+        "--no-jit-opts", action="store_true", help="Do not generate JIT-compiled optimizations for numerical code."
+    )
     parser.add_argument("--staging-review", action="store_true", help="Upload optimizations to staging for review")
     parser.add_argument(
         "--verify-setup",

codeflash/code_utils/code_extractor.py

@@ -26,12 +26,117 @@ if TYPE_CHECKING:
     from codeflash.models.models import FunctionSource
 
 
+class GlobalFunctionCollector(cst.CSTVisitor):
+    """Collects all module-level function definitions (not inside classes or other functions)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.functions: dict[str, cst.FunctionDef] = {}
+        self.function_order: list[str] = []
+        self.scope_depth = 0
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
+        if self.scope_depth == 0:
+            # Module-level function
+            name = node.name.value
+            self.functions[name] = node
+            if name not in self.function_order:
+                self.function_order.append(name)
+        self.scope_depth += 1
+        return True
+
+    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: ARG002
+        self.scope_depth -= 1
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:  # noqa: ARG002
+        self.scope_depth += 1
+        return True
+
+    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
+        self.scope_depth -= 1
+
+
+class GlobalFunctionTransformer(cst.CSTTransformer):
+    """Transforms/adds module-level functions from the new file to the original file."""
+
+    def __init__(self, new_functions: dict[str, cst.FunctionDef], new_function_order: list[str]) -> None:
+        super().__init__()
+        self.new_functions = new_functions
+        self.new_function_order = new_function_order
+        self.processed_functions: set[str] = set()
+        self.scope_depth = 0
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:  # noqa: ARG002
+        self.scope_depth += 1
+
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
+        self.scope_depth -= 1
+        if self.scope_depth > 0:
+            return updated_node
+
+        # Check if this is a module-level function we need to replace
+        name = original_node.name.value
+        if name in self.new_functions:
+            self.processed_functions.add(name)
+            return self.new_functions[name]
+        return updated_node
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> None:  # noqa: ARG002
+        self.scope_depth += 1
+
+    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:  # noqa: ARG002
+        self.scope_depth -= 1
+        return updated_node
+
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
+        # Add any new functions that weren't in the original file
+        new_statements = list(updated_node.body)
+
+        functions_to_append = [
+            self.new_functions[name]
+            for name in self.new_function_order
+            if name not in self.processed_functions and name in self.new_functions
+        ]
+
+        if functions_to_append:
+            # Find the position of the last function or class definition
+            insert_index = find_insertion_index_after_imports(updated_node)
+            for i, stmt in enumerate(new_statements):
+                if isinstance(stmt, (cst.FunctionDef, cst.ClassDef)):
+                    insert_index = i + 1
+
+            # Add empty line before each new function
+            function_nodes = []
+            for func in functions_to_append:
+                func_with_empty_line = func.with_changes(leading_lines=[cst.EmptyLine(), *func.leading_lines])
+                function_nodes.append(func_with_empty_line)
+
+            new_statements = list(chain(new_statements[:insert_index], function_nodes, new_statements[insert_index:]))
+
+        return updated_node.with_changes(body=new_statements)
+
+
+def collect_referenced_names(node: cst.CSTNode) -> set[str]:
+    """Collect all names referenced in a CST node using recursive traversal."""
+    names: set[str] = set()
+
+    def _collect(n: cst.CSTNode) -> None:
+        if isinstance(n, cst.Name):
+            names.add(n.value)
+        # Recursively process all children
+        for child in n.children:
+            _collect(child)
+
+    _collect(node)
+    return names
+
+
 class GlobalAssignmentCollector(cst.CSTVisitor):
     """Collects all global assignment statements."""
 
     def __init__(self) -> None:
         super().__init__()
-        self.assignments: dict[str, cst.Assign] = {}
+        self.assignments: dict[str, cst.Assign | cst.AnnAssign] = {}
         self.assignment_order: list[str] = []
         # Track scope depth to identify global assignments
         self.scope_depth = 0
@@ -73,6 +178,21 @@ class GlobalAssignmentCollector(cst.CSTVisitor):
                         self.assignment_order.append(name)
         return True
 
+    def visit_AnnAssign(self, node: cst.AnnAssign) -> Optional[bool]:
+        # Handle annotated assignments like: _CACHE: Dict[str, int] = {}
+        # Only process module-level annotated assignments with a value
+        if (
+            self.scope_depth == 0
+            and self.if_else_depth == 0
+            and isinstance(node.target, cst.Name)
+            and node.value is not None
+        ):
+            name = node.target.value
+            self.assignments[name] = node
+            if name not in self.assignment_order:
+                self.assignment_order.append(name)
+        return True
+
 
 def find_insertion_index_after_imports(node: cst.Module) -> int:
     """Find the position of the last import statement in the top-level of the module."""
@@ -104,7 +224,7 @@ def find_insertion_index_after_imports(node: cst.Module) -> int:
 class GlobalAssignmentTransformer(cst.CSTTransformer):
     """Transforms global assignments in the original file with those from the new file."""
 
-    def __init__(self, new_assignments: dict[str, cst.Assign], new_assignment_order: list[str]) -> None:
+    def __init__(self, new_assignments: dict[str, cst.Assign | cst.AnnAssign], new_assignment_order: list[str]) -> None:
         super().__init__()
         self.new_assignments = new_assignments
         self.new_assignment_order = new_assignment_order
@@ -151,38 +271,120 @@ class GlobalAssignmentTransformer(cst.CSTTransformer):
 
         return updated_node
 
+    def leave_AnnAssign(self, original_node: cst.AnnAssign, updated_node: cst.AnnAssign) -> cst.CSTNode:
+        if self.scope_depth > 0 or self.if_else_depth > 0:
+            return updated_node
+
+        # Check if this is a global annotated assignment we need to replace
+        if isinstance(original_node.target, cst.Name):
+            name = original_node.target.value
+            if name in self.new_assignments:
+                self.processed_assignments.add(name)
+                return self.new_assignments[name]
+
+        return updated_node
+
     def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
         # Add any new assignments that weren't in the original file
         new_statements = list(updated_node.body)
 
         # Find assignments to append
         assignments_to_append = [
-            self.new_assignments[name]
+            (name, self.new_assignments[name])
             for name in self.new_assignment_order
             if name not in self.processed_assignments and name in self.new_assignments
         ]
 
-        if assignments_to_append:
+        if not assignments_to_append:
-            # after last top-level imports
+            return updated_node.with_changes(body=new_statements)
+
+        # Collect all class and function names defined in the module
+        # These are the names that assignments might reference
+        module_defined_names: set[str] = set()
+        for stmt in new_statements:
+            if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)):
+                module_defined_names.add(stmt.name.value)
+
+        # Partition assignments: those that reference module definitions go at the end,
+        # those that don't can go right after imports
+        assignments_after_imports: list[tuple[str, cst.Assign | cst.AnnAssign]] = []
+        assignments_after_definitions: list[tuple[str, cst.Assign | cst.AnnAssign]] = []
+
+        for name, assignment in assignments_to_append:
+            # Get the value being assigned
+            if isinstance(assignment, (cst.Assign, cst.AnnAssign)) and assignment.value is not None:
+                value_node = assignment.value
+            else:
+                # No value to analyze, safe to place after imports
+                assignments_after_imports.append((name, assignment))
+                continue
+
+            # Collect names referenced in the assignment value
+            referenced_names = collect_referenced_names(value_node)
+
+            # Check if any referenced names are module-level definitions
+            if referenced_names & module_defined_names:
+                # This assignment references a class/function, place it after definitions
+                assignments_after_definitions.append((name, assignment))
+            else:
+                # Safe to place right after imports
+                assignments_after_imports.append((name, assignment))
+
+        # Insert assignments that don't depend on module definitions right after imports
+        if assignments_after_imports:
             insert_index = find_insertion_index_after_imports(updated_node)
+            assignment_lines = [
+                cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()])
+                for _, assignment in assignments_after_imports
+            ]
+            new_statements = list(chain(new_statements[:insert_index], assignment_lines, new_statements[insert_index:]))
+
+        # Insert assignments that depend on module definitions after all class/function definitions
+        if assignments_after_definitions:
+            # Find the position after the last function or class definition
+            insert_index = find_insertion_index_after_imports(cst.Module(body=new_statements))
+            for i, stmt in enumerate(new_statements):
+                if isinstance(stmt, (cst.FunctionDef, cst.ClassDef)):
+                    insert_index = i + 1
 
             assignment_lines = [
                 cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()])
-                for assignment in assignments_to_append
+                for _, assignment in assignments_after_definitions
             ]
-
             new_statements = list(chain(new_statements[:insert_index], assignment_lines, new_statements[insert_index:]))
 
-            # Add a blank line after the last assignment if needed
+        return updated_node.with_changes(body=new_statements)
-            after_index = insert_index + len(assignment_lines)
+
-            if after_index < len(new_statements):
+
-                next_stmt = new_statements[after_index]
+class GlobalStatementTransformer(cst.CSTTransformer):
-                # If there's no empty line, add one
+    """Transformer that appends global statements at the end of the module.
-                has_empty = any(isinstance(line, cst.EmptyLine) for line in next_stmt.leading_lines)
+
-                if not has_empty:
+    This ensures that global statements (like function calls at module level) are placed
-                    new_statements[after_index] = next_stmt.with_changes(
+    after all functions, classes, and assignments they might reference, preventing NameError
-                        leading_lines=[cst.EmptyLine(), *next_stmt.leading_lines]
+    at module load time.
-                    )
+
+    This transformer should be run LAST after GlobalFunctionTransformer and
+    GlobalAssignmentTransformer have already added their content.
+    """
+
+    def __init__(self, global_statements: list[cst.SimpleStatementLine]) -> None:
+        super().__init__()
+        self.global_statements = global_statements
+
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
+        if not self.global_statements:
+            return updated_node
+
+        new_statements = list(updated_node.body)
+
+        # Add empty line before each statement for readability
+        statement_lines = [
+            stmt.with_changes(leading_lines=[cst.EmptyLine(), *stmt.leading_lines]) for stmt in self.global_statements
+        ]
+
+        # Append statements at the end of the module
+        # This ensures they come after all functions, classes, and assignments
+        new_statements.extend(statement_lines)
 
         return updated_node.with_changes(body=new_statements)
 
@@ -214,8 +416,8 @@ class GlobalStatementCollector(cst.CSTVisitor):
     def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
         if not self.in_function_or_class:
             for statement in node.body:
-                # Skip imports
+                # Skip imports and assignments (both regular and annotated)
-                if not isinstance(statement, (cst.Import, cst.ImportFrom, cst.Assign)):
+                if not isinstance(statement, (cst.Import, cst.ImportFrom, cst.Assign, cst.AnnAssign)):
                     self.global_statements.append(node)
                     break
 
@@ -310,40 +512,6 @@ class DottedImportCollector(cst.CSTVisitor):
             self._collect_imports_from_block(node.body)
 
 
-class ImportInserter(cst.CSTTransformer):
-    """Transformer that inserts global statements after the last import."""
-
-    def __init__(self, global_statements: list[cst.SimpleStatementLine], last_import_line: int) -> None:
-        super().__init__()
-        self.global_statements = global_statements
-        self.last_import_line = last_import_line
-        self.current_line = 0
-        self.inserted = False
-
-    def leave_SimpleStatementLine(
-        self,
-        original_node: cst.SimpleStatementLine,  # noqa: ARG002
-        updated_node: cst.SimpleStatementLine,
-    ) -> cst.Module:
-        self.current_line += 1
-
-        # If we're right after the last import and haven't inserted yet
-        if self.current_line == self.last_import_line and not self.inserted:
-            self.inserted = True
-            return cst.Module(body=[updated_node, *self.global_statements])
-
-        return cst.Module(body=[updated_node])
-
-    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
-        # If there were no imports, add at the beginning of the module
-        if self.last_import_line == 0 and not self.inserted:
-            updated_body = list(updated_node.body)
-            for stmt in reversed(self.global_statements):
-                updated_body.insert(0, stmt)
-            return updated_node.with_changes(body=updated_body)
-        return updated_node
-
-
 def extract_global_statements(source_code: str) -> tuple[cst.Module, list[cst.SimpleStatementLine]]:
     """Extract global statements from source code."""
     module = cst.parse_module(source_code)
@@ -395,34 +563,58 @@ def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
             continue
         unique_global_statements.append(stmt)
 
-    mod_dst_code = dst_module_code
+    # Reuse already-parsed dst_module
-    # Insert unique global statements if any
+    original_module = dst_module
-    if unique_global_statements:
-        last_import_line = find_last_import_line(dst_module_code)
-        # Reuse already-parsed dst_module
-        transformer = ImportInserter(unique_global_statements, last_import_line)
-        # Use visit inplace, don't parse again
-        modified_module = dst_module.visit(transformer)
-        mod_dst_code = modified_module.code
-        # Parse the code after insertion
-        original_module = cst.parse_module(mod_dst_code)
-    else:
-        # No new statements to insert, reuse already-parsed dst_module
-        original_module = dst_module
 
     # Parse the src_module_code once only (already done above: src_module)
     # Collect assignments from the new file
-    new_collector = GlobalAssignmentCollector()
+    new_assignment_collector = GlobalAssignmentCollector()
-    src_module.visit(new_collector)
+    src_module.visit(new_assignment_collector)
-    # Only create transformer if there are assignments to insert/transform
-    if not new_collector.assignments:  # nothing to transform
-        return mod_dst_code
 
-    # Transform the original destination module
+    # Collect module-level functions from both source and destination
-    transformer = GlobalAssignmentTransformer(new_collector.assignments, new_collector.assignment_order)
+    src_function_collector = GlobalFunctionCollector()
-    transformed_module = original_module.visit(transformer)
+    src_module.visit(src_function_collector)
 
-    return transformed_module.code
+    dst_function_collector = GlobalFunctionCollector()
+    original_module.visit(dst_function_collector)
+
+    # Filter out functions that already exist in the destination (only add truly new functions)
+    new_functions = {
+        name: func
+        for name, func in src_function_collector.functions.items()
+        if name not in dst_function_collector.functions
+    }
+    new_function_order = [name for name in src_function_collector.function_order if name in new_functions]
+
+    # If there are no assignments, no new functions, and no global statements, return unchanged
+    if not new_assignment_collector.assignments and not new_functions and not unique_global_statements:
+        return dst_module_code
+
+    # The order of transformations matters:
+    # 1. Functions first - so assignments and statements can reference them
+    # 2. Assignments second - so they come after functions but before statements
+    # 3. Global statements last - so they can reference both functions and assignments
+
+    # Transform functions if any
+    if new_functions:
+        function_transformer = GlobalFunctionTransformer(new_functions, new_function_order)
+        original_module = original_module.visit(function_transformer)
+
+    # Transform assignments if any
+    if new_assignment_collector.assignments:
+        transformer = GlobalAssignmentTransformer(
+            new_assignment_collector.assignments, new_assignment_collector.assignment_order
+        )
+        original_module = original_module.visit(transformer)
+
+    # Insert global statements (like function calls at module level) LAST,
+    # after all functions and assignments are added, to ensure they can reference any
+    # functions or variables defined in the module
+    if unique_global_statements:
+        statement_transformer = GlobalStatementTransformer(unique_global_statements)
+        original_module = original_module.visit(statement_transformer)
+
+    return original_module.code
 
 
 def resolve_star_import(module_name: str, project_root: Path) -> set[str]:

codeflash/context/unused_definition_remover.py

@@ -295,11 +295,18 @@ class DependencyCollector(cst.CSTVisitor):
             return
 
         if name in self.definitions and name != self.current_top_level_name:
-            # skip if we are refrencing a class attribute and not a top-level definition
+            # Skip if this Name is the .attr part of an Attribute (e.g., 'x' in 'self.x')
+            # We only want to track the base/value of attribute access, not the attribute name itself
             if self.class_depth > 0:
                 parent = self.get_metadata(cst.metadata.ParentNodeProvider, node)
                 if parent is not None and isinstance(parent, cst.Attribute):
-                    return
+                    # Check if this Name is the .attr (property name), not the .value (base)
+                    # If it's the .attr, skip it - attribute names aren't references to definitions
+                    if parent.attr is node:
+                        return
+                    # If it's the .value (base), only skip if it's self/cls
+                    if name in ("self", "cls"):
+                        return
             self.definitions[self.current_top_level_name].dependencies.add(name)
 
 
@@ -553,16 +560,6 @@ def remove_unused_definitions_by_function_names(code: str, qualified_function_na
         return code
 
 
-def print_definitions(definitions: dict[str, UsageInfo]) -> None:
-    """Print information about each definition without the complex node object, used for debugging."""
-    print(f"Found {len(definitions)} definitions:")
-    for name, info in sorted(definitions.items()):
-        print(f"  - Name: {name}")
-        print(f"    Used by qualified function: {info.used_by_qualified_function}")
-        print(f"    Dependencies: {', '.join(sorted(info.dependencies)) if info.dependencies else 'None'}")
-        print()
-
-
 def revert_unused_helper_functions(
     project_root: Path, unused_helpers: list[FunctionSource], original_helper_code: dict[Path, str]
 ) -> None:
@@ -637,43 +634,40 @@ def _analyze_imports_in_optimized_code(
             func_name = helper.only_function_name
             module_name = helper.file_path.stem
             # Cache function lookup for this (module, func)
-            file_entry = helpers_by_file_and_func[module_name]
+            helpers_by_file_and_func[module_name].setdefault(func_name, []).append(helper)
-            if func_name in file_entry:
-                file_entry[func_name].append(helper)
-            else:
-                file_entry[func_name] = [helper]
             helpers_by_file[module_name].append(helper)
 
-    # Optimize attribute lookups and method binding outside the loop
-    helpers_by_file_and_func_get = helpers_by_file_and_func.get
-    helpers_by_file_get = helpers_by_file.get
-
     for node in ast.walk(optimized_ast):
         if isinstance(node, ast.ImportFrom):
             # Handle "from module import function" statements
             module_name = node.module
             if module_name:
-                file_entry = helpers_by_file_and_func_get(module_name, None)
+                file_entry = helpers_by_file_and_func.get(module_name)
                 if file_entry:
                     for alias in node.names:
                         imported_name = alias.asname if alias.asname else alias.name
                         original_name = alias.name
-                        helpers = file_entry.get(original_name, None)
+                        helpers = file_entry.get(original_name)
                         if helpers:
+                            imported_set = imported_names_map[imported_name]
                             for helper in helpers:
-                                imported_names_map[imported_name].add(helper.qualified_name)
+                                imported_set.add(helper.qualified_name)
-                                imported_names_map[imported_name].add(helper.fully_qualified_name)
+                                imported_set.add(helper.fully_qualified_name)
 
         elif isinstance(node, ast.Import):
             # Handle "import module" statements
             for alias in node.names:
                 imported_name = alias.asname if alias.asname else alias.name
                 module_name = alias.name
-                for helper in helpers_by_file_get(module_name, []):
+                helpers = helpers_by_file.get(module_name)
-                    # For "import module" statements, functions would be called as module.function
+                if helpers:
-                    full_call = f"{imported_name}.{helper.only_function_name}"
+                    imported_set = imported_names_map[f"{imported_name}.{{func}}"]
-                    imported_names_map[full_call].add(helper.qualified_name)
+                    for helper in helpers:
-                    imported_names_map[full_call].add(helper.fully_qualified_name)
+                        # For "import module" statements, functions would be called as module.function
+                        full_call = f"{imported_name}.{helper.only_function_name}"
+                        full_call_set = imported_names_map[full_call]
+                        full_call_set.add(helper.qualified_name)
+                        full_call_set.add(helper.fully_qualified_name)
 
     return dict(imported_names_map)
 
@@ -758,27 +752,31 @@ def detect_unused_helper_functions(
                     called_name = node.func.id
                     called_function_names.add(called_name)
                     # Also add the qualified name if this is an imported function
-                    if called_name in imported_names_map:
+                    mapped_names = imported_names_map.get(called_name)
-                        called_function_names.update(imported_names_map[called_name])
+                    if mapped_names:
+                        called_function_names.update(mapped_names)
                 elif isinstance(node.func, ast.Attribute):
                     # Method call: obj.method() or self.method() or module.function()
                     if isinstance(node.func.value, ast.Name):
-                        if node.func.value.id == "self":
+                        attr_name = node.func.attr
+                        value_id = node.func.value.id
+                        if value_id == "self":
                             # self.method_name() -> add both method_name and ClassName.method_name
-                            called_function_names.add(node.func.attr)
+                            called_function_names.add(attr_name)
+                            # For class methods, also add the qualified name
                             # For class methods, also add the qualified name
                             if hasattr(function_to_optimize, "parents") and function_to_optimize.parents:
                                 class_name = function_to_optimize.parents[0].name
-                                called_function_names.add(f"{class_name}.{node.func.attr}")
+                                called_function_names.add(f"{class_name}.{attr_name}")
                         else:
-                            # obj.method() or module.function()
-                            attr_name = node.func.attr
                             called_function_names.add(attr_name)
-                            called_function_names.add(f"{node.func.value.id}.{attr_name}")
+                            full_call = f"{value_id}.{attr_name}"
+                            called_function_names.add(full_call)
                             # Check if this is a module.function call that maps to a helper
-                            full_call = f"{node.func.value.id}.{attr_name}"
+                            mapped_names = imported_names_map.get(full_call)
-                            if full_call in imported_names_map:
+                            if mapped_names:
-                                called_function_names.update(imported_names_map[full_call])
+                                called_function_names.update(mapped_names)
+                    # Handle nested attribute access like obj.attr.method()
                     # Handle nested attribute access like obj.attr.method()
                     else:
                         called_function_names.add(node.func.attr)
@@ -788,6 +786,7 @@ def detect_unused_helper_functions(
 
         # Find helper functions that are no longer called
         unused_helpers = []
+        entrypoint_file_path = function_to_optimize.file_path
         for helper_function in code_context.helper_functions:
             jedi_type = helper_function.jedi_definition.type if helper_function.jedi_definition else None
             if jedi_type != "class":  # Include when jedi_definition is None (non-Python)
@@ -796,29 +795,30 @@ def detect_unused_helper_functions(
                 helper_simple_name = helper_function.only_function_name
                 helper_fully_qualified_name = helper_function.fully_qualified_name
 
-                # Create a set of all possible names this helper might be called by
+                # Check membership efficiently - exit early on first match
-                possible_call_names = {helper_qualified_name, helper_simple_name, helper_fully_qualified_name}
+                if (
-
+                    helper_qualified_name in called_function_names
+                    or helper_simple_name in called_function_names
+                    or helper_fully_qualified_name in called_function_names
+                ):
+                    is_called = True
                 # For cross-file helpers, also consider module-based calls
-                if helper_function.file_path != function_to_optimize.file_path:
+                elif helper_function.file_path != entrypoint_file_path:
                     # Add potential module.function combinations
                     module_name = helper_function.file_path.stem
-                    possible_call_names.add(f"{module_name}.{helper_simple_name}")
+                    module_call = f"{module_name}.{helper_simple_name}"
-
+                    is_called = module_call in called_function_names
-                # Check if any of the possible names are in the called functions
+                else:
-                is_called = bool(possible_call_names.intersection(called_function_names))
+                    is_called = False
 
                 if not is_called:
                     unused_helpers.append(helper_function)
                     logger.debug(f"Helper function {helper_qualified_name} is not called in optimized code")
-                    logger.debug(f"  Checked names: {possible_call_names}")
                 else:
                     logger.debug(f"Helper function {helper_qualified_name} is still called in optimized code")
-                    logger.debug(f"  Called via: {possible_call_names.intersection(called_function_names)}")
-
-        ret_val = unused_helpers
 
     except Exception as e:
         logger.debug(f"Error detecting unused helper functions: {e}")
-        ret_val = []
+        return []
-    return ret_val
+    else:
+        return unused_helpers

codeflash/optimization/function_optimizer.py

@@ -700,7 +700,9 @@ class FunctionOptimizer:
         ):
             console.rule()
             new_code_context = code_context
-            if self.is_numerical_code:  # if the code is numerical in nature (uses numpy/tensorflow/math/pytorch/jax)
+            if (
+                self.is_numerical_code and not self.args.no_jit_opts
+            ):  # if the code is numerical in nature (uses numpy/tensorflow/math/pytorch/jax)
                 jit_compiled_opt_candidate = self.aiservice_client.get_jit_rewritten_code(
                     code_context.read_writable_code.markdown, self.function_trace_id
                 )
@@ -729,7 +731,7 @@ class FunctionOptimizer:
                 read_writable_code=code_context.read_writable_code,
                 read_only_context_code=code_context.read_only_context_code,
                 run_experiment=should_run_experiment,
-                is_numerical_code=self.is_numerical_code,
+                is_numerical_code=self.is_numerical_code and not self.args.no_jit_opts,
             )
 
             concurrent.futures.wait([future_tests, future_optimizations])
@@ -1251,7 +1253,7 @@ class FunctionOptimizer:
             )
             if self.experiment_id
             else None,
-            is_numerical_code=self.is_numerical_code,
+            is_numerical_code=self.is_numerical_code and not self.args.no_jit_opts,
             language=self.function_to_optimize.language,
         )
 

codeflash/tracer.py

@@ -138,6 +138,13 @@ def main(args: Namespace | None = None) -> ArgumentParser:
                         env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}"
                     else:
                         env["PYTHONPATH"] = project_root_str
+                    # Disable JIT compilation to ensure tracing captures all function calls
+                    env["NUMBA_DISABLE_JIT"] = str(1)
+                    env["TORCHDYNAMO_DISABLE"] = str(1)
+                    env["PYTORCH_JIT"] = str(0)
+                    env["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=0"
+                    env["TF_ENABLE_ONEDNN_OPTS"] = str(0)
+                    env["JAX_DISABLE_JIT"] = str(1)
                     processes.append(
                         subprocess.Popen(
                             [
@@ -175,6 +182,13 @@ def main(args: Namespace | None = None) -> ArgumentParser:
                     env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}"
                 else:
                     env["PYTHONPATH"] = project_root_str
+                # Disable JIT compilation to ensure tracing captures all function calls
+                env["NUMBA_DISABLE_JIT"] = str(1)
+                env["TORCHDYNAMO_DISABLE"] = str(1)
+                env["PYTORCH_JIT"] = str(0)
+                env["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=0"
+                env["TF_ENABLE_ONEDNN_OPTS"] = str(0)
+                env["JAX_DISABLE_JIT"] = str(1)
 
                 subprocess.run(
                     [

codeflash/verification/codeflash_capture.py

@@ -15,6 +15,8 @@ from typing import Callable
 import dill as pickle
 from dill import PicklingWarning
 
+from codeflash.picklepatch.pickle_patcher import PicklePatcher
+
 warnings.filterwarnings("ignore", category=PicklingWarning)
 
 
@@ -148,18 +150,29 @@ def codeflash_capture(function_name: str, tmp_dir_path: str, tests_root: str, is
             print(f"!######{test_stdout_tag}######!")
 
             # Capture instance state after initialization
-            if hasattr(args[0], "__dict__"):
+            # self is always the first argument, this is ensured during instrumentation
-                instance_state = args[
+            instance = args[0]
-                    0
+            if hasattr(instance, "__dict__"):
-                ].__dict__  # self is always the first argument, this is ensured during instrumentation
+                instance_state = instance.__dict__
+            elif hasattr(instance, "__slots__"):
+                # For classes using __slots__, capture slot values
+                instance_state = {
+                    slot: getattr(instance, slot, None) for slot in instance.__slots__ if hasattr(instance, slot)
+                }
             else:
-                raise ValueError("Instance state could not be captured.")
+                # For C extension types or other special classes (e.g., Playwright's Page),
+                # capture all non-private, non-callable attributes
+                instance_state = {
+                    attr: getattr(instance, attr)
+                    for attr in dir(instance)
+                    if not attr.startswith("_") and not callable(getattr(instance, attr, None))
+                }
             codeflash_cur.execute(
                 "CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT)"
             )
 
             # Write to sqlite
-            pickled_return_value = pickle.dumps(exception) if exception else pickle.dumps(instance_state)
+            pickled_return_value = pickle.dumps(exception) if exception else PicklePatcher.dumps(instance_state)
             codeflash_cur.execute(
                 "INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                 (

codeflash/verification/comparator.py

@@ -257,6 +257,10 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool:  # noqa: ANN001
                     return False
             return True
 
+        # Handle mappingproxy (read-only dict view, commonly seen as class.__dict__)
+        if isinstance(orig, types.MappingProxyType):
+            return comparator(dict(orig), dict(new), superset_obj)
+
         # Handle dict view types (dict_keys, dict_values, dict_items)
         # Use type name checking since these are not directly importable types
         type_name = type(orig).__name__

codeflash/verification/equivalence.py

@@ -18,6 +18,12 @@ reprlib_repr = reprlib.Repr()
 reprlib_repr.maxstring = 1500
 test_diff_repr = reprlib_repr.repr
 
+def safe_repr(obj: object) -> str:
+    """Safely get repr of an object, handling Mock objects with corrupted state."""
+    try:
+        return repr(obj)
+    except (AttributeError, TypeError, RecursionError) as e:
+        return f"<repr failed: {type(e).__name__}: {e}>"
 
 def compare_test_results(
     original_results: TestResults,
@@ -102,8 +108,8 @@ def compare_test_results(
             test_diffs.append(
                 TestDiff(
                     scope=TestDiffScope.RETURN_VALUE,
-                    original_value=test_diff_repr(repr(original_test_result.return_value)),
+                    original_value=test_diff_repr(safe_repr(original_test_result.return_value)),
-                    candidate_value=test_diff_repr(repr(cdd_test_result.return_value)),
+                    candidate_value=test_diff_repr(safe_repr(cdd_test_result.return_value)),
                     test_src_code=original_test_result.id.get_src_code(original_test_result.file_name),
                     candidate_pytest_error=cdd_pytest_error,
                     original_pass=original_test_result.did_pass,

codeflash/version.py

@@ -1,2 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.19.1.post300.dev0+ad95a41b"
+__version__ = "0.20.0"

docs/cli-reference/flags.mdx

@@ -1,191 +0,0 @@
----
-title: "Flags Reference"
-description: "Complete reference for all Codeflash CLI flags and options"
-icon: "list"
-sidebarTitle: "Flags Reference"
-keywords: ["flags", "options", "arguments", "command line"]
----
-
-# Flags Reference
-
-Complete reference for all Codeflash CLI flags and command-line options.
-
----
-
-## Main Command Flags
-
-| Flag | Type | Description |
-|------|------|-------------|
-| `--file` | `PATH` | Optimize only this file |
-| `--function` | `NAME` | Optimize only this function (requires `--file`) |
-| `--all` | `[PATH]` | Optimize all functions. Optional path to start from |
-| `--replay-test` | `PATH` | Path to replay test file(s) |
-| `--benchmark` | flag | Enable benchmark mode |
-| `--no-pr` | flag | Don't create PR, update locally |
-| `--no-gen-tests` | flag | Don't generate tests |
-| `--no-draft` | flag | Skip draft PRs |
-| `--worktree` | flag | Use git worktree |
-| `--staging-review` | flag | Upload to staging |
-| `--verbose` / `-v` | flag | Verbose debug output |
-| `--verify-setup` | flag | Run setup verification |
-| `--version` | flag | Show version |
-
----
-
-## Configuration Override Flags
-
-Override settings from `pyproject.toml` via command line.
-
-| Flag | Type | Description |
-|------|------|-------------|
-| `--config-file` | `PATH` | Path to pyproject.toml |
-| `--module-root` | `PATH` | Python module root directory |
-| `--tests-root` | `PATH` | Tests directory |
-| `--benchmarks-root` | `PATH` | Benchmarks directory |
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Override config file location
-    codeflash --file src/app.py --function main --config-file configs/pyproject.toml --no-pr
-
-    # Override module root
-    codeflash --file src/app.py --function main --module-root src --no-pr
-
-    # Override tests root
-    codeflash --file src/app.py --function main --tests-root tests/unit --no-pr
-
-    # Combine multiple overrides
-    codeflash --file src/app.py --function main \
-      --module-root src \
-      --tests-root tests \
-      --no-pr
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Override config file location
-    codeflash --file src\app.py --function main --config-file configs\pyproject.toml --no-pr
-
-    # Override module root
-    codeflash --file src\app.py --function main --module-root src --no-pr
-
-    # Override tests root
-    codeflash --file src\app.py --function main --tests-root tests\unit --no-pr
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
----
-
-## Optimize Subcommand Flags
-
-Flags specific to the `codeflash optimize` command.
-
-| Flag | Type | Description |
-|------|------|-------------|
-| `--output` | `PATH` | Trace file output path (default: `codeflash.trace`) |
-| `--timeout` | `INT` | Maximum trace time in seconds |
-| `--max-function-count` | `INT` | Max times to trace a function (default: 100) |
-| `--config-file-path` | `PATH` | Path to pyproject.toml |
-| `--trace-only` | flag | Only trace, don't optimize |
-
-<Info>
-  The `--output` flag specifies where to save the trace file. If not specified, it defaults to `codeflash.trace` in the current directory.
-</Info>
-
----
-
-## Behavior Flags
-
-Control how Codeflash behaves during optimization.
-
-| Flag | Description |
-|------|-------------|
-| `--no-pr` | Run locally without creating a pull request |
-| `--no-gen-tests` | Use only existing tests, skip test generation |
-| `--no-draft` | Skip optimization for draft PRs (CI mode) |
-| `--worktree` | Use git worktree for isolated optimization |
-| `--staging-review` | Upload optimizations to staging for review |
-| `--verbose` / `-v` | Enable verbose debug logging |
-
-<Accordion title="Complete Examples">
-```bash
-# Local optimization only
-codeflash --file src/app.py --function main --no-pr
-
-# Use only existing tests
-codeflash --file src/app.py --function main --no-gen-tests --no-pr
-
-# Enable verbose logging
-codeflash --file src/app.py --function main --verbose --no-pr
-
-# Use worktree for isolation
-codeflash --file src/app.py --function main --worktree --no-pr
-
-# Upload to staging
-codeflash --all --staging-review --no-pr
-```
-</Accordion>
-
----
-
-## Flag Combinations
-
-Common flag combinations for different use cases:
-
-### Local Development
-
-```bash
-# Optimize locally with verbose output
-codeflash --file src/app.py --function main --no-pr --verbose
-```
-
-### CI/CD Pipeline
-
-```bash
-# Skip draft PRs and use existing tests only
-codeflash --all --no-draft --no-gen-tests
-```
-
-### Debugging
-
-```bash
-# Trace only with custom output and timeout
-codeflash optimize app.py --trace-only --output debug.trace --timeout 60
-```
-
-### Custom Configuration
-
-```bash
-# Override multiple config settings
-codeflash --file src/app.py --function main \
-  --module-root src \
-  --tests-root tests/unit \
-  --benchmarks-root tests/benchmarks \
-  --no-pr
-```
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Optimization Commands"
-    icon="bullseye"
-    href="/cli-reference/optimization"
-  >
-    Learn how to use optimization commands
-  </Card>
-  <Card
-    title="Troubleshooting"
-    icon="wrench"
-    href="/cli-reference/troubleshooting"
-  >
-    Fix common issues
-  </Card>
-</CardGroup>
-

docs/cli-reference/index.mdx

@@ -1,208 +0,0 @@
----
-title: "CLI Reference"
-description: "Complete command-line reference for Codeflash CLI commands, flags, and options"
-icon: "terminal"
-sidebarTitle: "Overview"
-keywords:
-  [
-    "CLI",
-    "command line",
-    "commands",
-    "flags",
-    "options",
-    "reference",
-    "terminal",
-  ]
----
-
-# Codeflash CLI Reference
-
-Complete command-line reference for all Codeflash commands, flags, and options with practical examples you can run directly in your terminal.
-
-<Info>
-  **Prerequisites** - Ensure Codeflash is installed in your Python environment
-  and you have a configured `pyproject.toml` in your project.
-</Info>
-
----
-
-## Quick Start
-
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Activate virtual environment (if using one)
-    source .venv/bin/activate
-
-    # Verify installation
-    codeflash --version
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Activate virtual environment (if using one)
-    .venv\Scripts\activate
-
-    # Verify installation
-    codeflash --version
-    ```
-  </Tab>
-</Tabs>
-
----
-
-## Common Workflows
-
-### 1. First-Time Setup
-
-<Steps>
-  <Step title="Install Codeflash">
-    ```bash
-    pip install codeflash
-    ```
-  </Step>
-  <Step title="Initialize Project">
-    ```bash
-    codeflash init
-    ```
-  </Step>
-  <Step title="Verify Setup">
-    ```bash
-    codeflash --verify-setup
-    ```
-  </Step>
-  <Step title="Run First Optimization">
-    ```bash
-    codeflash --file src/main.py --function my_function --no-pr
-    ```
-  </Step>
-</Steps>
-
----
-
-### 2. Optimize a Workflow
-
-<Steps>
-  <Step title="Trace Your Script">
-    ```bash
-    codeflash optimize my_script.py --arg1 value1
-    ```
-  </Step>
-  <Step title="Review Optimizations">
-    Check the generated PR or local changes for optimization suggestions.
-  </Step>
-</Steps>
-
----
-
-### 3. CI/CD Integration
-
-<Steps>
-  <Step title="Set Up GitHub Actions">
-    ```bash
-    codeflash init-actions
-    ```
-  </Step>
-  <Step title="Merge the Workflow PR">
-    Review and merge the generated GitHub Actions workflow.
-  </Step>
-  <Step title="Automatic Optimization">
-    Codeflash will now optimize code in every PR automatically!
-  </Step>
-</Steps>
-
----
-
-## Help & Version
-
-```bash
-# Display version
-codeflash --version
-
-# Main help
-codeflash --help
-
-# Subcommand help
-codeflash optimize --help
-codeflash init --help
-```
-
----
-
-## Documentation Structure
-
-This CLI reference is organized into the following sections:
-
-<CardGroup cols={2}>
-  <Card
-    title="Setup Commands"
-    icon="wrench"
-    href="/cli-reference/setup"
-  >
-    Initialize projects, set up GitHub Actions, and verify installation
-  </Card>
-  <Card
-    title="Optimization Commands"
-    icon="bullseye"
-    href="/cli-reference/optimization"
-  >
-    Optimize single functions or entire codebases
-  </Card>
-  <Card
-    title="Tracing & Workflows"
-    icon="route"
-    href="/cli-reference/tracing"
-  >
-    Trace script execution and optimize based on real usage
-  </Card>
-  <Card
-    title="Flags Reference"
-    icon="list"
-    href="/cli-reference/flags"
-  >
-    Complete reference for all command-line flags
-  </Card>
-  <Card
-    title="Troubleshooting"
-    icon="wrench"
-    href="/cli-reference/troubleshooting"
-  >
-    Solutions for common CLI issues
-  </Card>
-</CardGroup>
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Optimize a Function"
-    icon="bullseye"
-    href="/optimizing-with-codeflash/one-function"
-  >
-    Learn how to optimize individual functions
-  </Card>
-  <Card
-    title="Trace & Optimize"
-    icon="route"
-    href="/optimizing-with-codeflash/trace-and-optimize"
-  >
-    Optimize entire workflows with tracing
-  </Card>
-  <Card
-    title="GitHub Actions"
-    icon="github"
-    href="/optimizing-with-codeflash/codeflash-github-actions"
-  >
-    Set up continuous optimization
-  </Card>
-  <Card
-    title="Configuration"
-    icon="gear"
-    href="/configuration"
-  >
-    Advanced configuration options
-  </Card>
-</CardGroup>
-

docs/cli-reference/optimization.mdx

@@ -1,172 +0,0 @@
----
-title: "Optimization Commands"
-description: "Optimize single functions or entire codebases with Codeflash CLI"
-icon: "bullseye"
-sidebarTitle: "Optimization Commands"
-keywords: ["optimization", "function", "file", "all", "commands"]
----
-
-# Optimization Commands
-
-Commands for optimizing individual functions or entire codebases.
-
----
-
-## Optimize a Single Function
-
-Target a specific function in a file for optimization.
-
-```bash
-codeflash --file <path/to/file.py> --function <function_name>
-```
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Basic optimization (creates PR)
-    codeflash --file src/utils.py --function calculate_metrics
-
-    # Local optimization only (no PR)
-    codeflash --file src/utils.py --function calculate_metrics --no-pr
-
-    # With verbose output
-    codeflash --file src/utils.py --function calculate_metrics --no-pr --verbose
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Basic optimization (creates PR)
-    codeflash --file src\utils.py --function calculate_metrics
-
-    # Local optimization only (no PR)
-    codeflash --file src\utils.py --function calculate_metrics --no-pr
-
-    # With verbose output
-    codeflash --file src\utils.py --function calculate_metrics --no-pr --verbose
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
-<Warning>
-  **Important**: The file must be within your configured `module-root`
-  directory. Files outside `module-root` will be ignored with "Functions outside
-  module-root" message.
-</Warning>
-
----
-
-## Optimize All Functions
-
-Optimize all functions in your entire codebase or a specific directory.
-
-```bash
-# Optimize entire codebase
-codeflash --all
-
-# Optimize specific directory
-codeflash --all src/core/
-```
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Optimize all (creates PRs)
-    codeflash --all
-
-    # Optimize all locally (no PRs)
-    codeflash --all --no-pr
-
-    # Optimize specific directory
-    codeflash --all src/algorithms/ --no-pr
-
-    # Skip draft PRs in CI
-    codeflash --all --no-draft
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Optimize all (creates PRs)
-    codeflash --all
-
-    # Optimize all locally (no PRs)
-    codeflash --all --no-pr
-
-    # Optimize specific directory
-    codeflash --all src\algorithms\ --no-pr
-
-    # Skip draft PRs in CI
-    codeflash --all --no-draft
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
-<Info>
-  When using `--all`, Codeflash will:
-  - Discover all optimizable functions in your codebase
-  - Create separate PRs for each function (or update locally with `--no-pr`)
-  - Process functions in batches to avoid overwhelming your repository
-</Info>
-
----
-
-## Benchmark Mode
-
-Optimize code based on performance benchmarks using pytest-benchmark format.
-
-```bash
-codeflash --file <file.py> --benchmark --benchmarks-root <path>
-```
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # With benchmarks-root flag
-    codeflash --file src/core.py --benchmark --benchmarks-root tests/benchmarks --no-pr
-
-    # If benchmarks-root is in pyproject.toml
-    codeflash --file src/core.py --benchmark --no-pr
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # With benchmarks-root flag
-    codeflash --file src\core.py --benchmark --benchmarks-root tests\benchmarks --no-pr
-
-    # If benchmarks-root is in pyproject.toml
-    codeflash --file src\core.py --benchmark --no-pr
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
-<Warning>
-  The `--benchmarks-root` directory must exist and be configured either via
-  `pyproject.toml` or the command-line flag.
-</Warning>
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Tracing & Workflows"
-    icon="route"
-    href="/cli-reference/tracing"
-  >
-    Learn about trace-based optimization
-  </Card>
-  <Card
-    title="Flags Reference"
-    icon="list"
-    href="/cli-reference/flags"
-  >
-    Complete flag reference
-  </Card>
-</CardGroup>
-

docs/cli-reference/setup.mdx

@@ -1,125 +0,0 @@
----
-title: "Setup Commands"
-description: "Initialize projects, set up GitHub Actions, and verify Codeflash installation"
-icon: "wrench"
-sidebarTitle: "Setup Commands"
-keywords: ["setup", "init", "installation", "github actions", "verify"]
----
-
-# Setup Commands
-
-Commands for initializing Codeflash in your project, setting up continuous optimization, and verifying your installation.
-
----
-
-## `codeflash init`
-
-Initialize Codeflash for your Python project. This creates the configuration in `pyproject.toml`.
-
-<CodeGroup>
-```bash Basic
-codeflash init
-```
-
-```bash With Formatter Override
-codeflash init --override-formatter-check
-```
-</CodeGroup>
-
-<Tip>
-  The `init` command will guide you through an interactive setup process,
-  including API key configuration, module selection, and GitHub App
-  installation.
-</Tip>
-
-**What it does:**
-
-- Prompts for your Python module directory (`module-root`)
-- Prompts for your test directory (`tests-root`)
-- Configures code formatter preferences
-- Sets up telemetry preferences
-- Optionally installs the Codeflash VS Code extension
-- Optionally sets up GitHub Actions workflow
-
----
-
-## `codeflash init-actions`
-
-Set up GitHub Actions workflow for continuous optimization on every pull request.
-
-```bash
-codeflash init-actions
-```
-
-**What it does:**
-
-- Creates a workflow file in `.github/workflows/`
-- Opens a PR with the workflow configuration
-- Requires the Codeflash GitHub App to be installed
-
-<Warning>
-  This command requires the Codeflash GitHub App to be installed on your repository. If you haven't installed it, you'll be prompted with a link to do so.
-</Warning>
-
----
-
-## `codeflash vscode-install`
-
-Install the Codeflash extension for VS Code, Cursor, or Windsurf.
-
-```bash
-codeflash vscode-install
-```
-
-**What it does:**
-
-- Detects which editor you're using (VS Code, Cursor, or Windsurf)
-- Downloads and installs the appropriate extension
-- Works with both Marketplace and Open VSX sources
-
-<Tip>
-  This command is also run automatically during `codeflash init` if you choose to install the extension.
-</Tip>
-
----
-
-## `codeflash --verify-setup`
-
-Verify your Codeflash installation by running a sample optimization.
-
-```bash
-codeflash --verify-setup
-```
-
-**What it does:**
-
-- Creates a temporary demo file
-- Runs a sample optimization
-- Verifies all components are working correctly
-- Cleans up the demo file afterward
-
-<Note>
-  This command takes about 3 minutes to complete. It's a great way to ensure everything is set up correctly before optimizing your actual code.
-</Note>
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Optimization Commands"
-    icon="bullseye"
-    href="/cli-reference/optimization"
-  >
-    Learn how to optimize functions
-  </Card>
-  <Card
-    title="Flags Reference"
-    icon="list"
-    href="/cli-reference/flags"
-  >
-    Complete flag reference
-  </Card>
-</CardGroup>
-

docs/cli-reference/tracing.mdx

@@ -1,213 +0,0 @@
----
-title: "Tracing & Workflows"
-description: "Trace script execution and optimize functions based on real-world usage"
-icon: "route"
-sidebarTitle: "Tracing & Workflows"
-keywords: ["tracing", "optimize", "workflow", "replay test", "pytest"]
----
-
-# Tracing & Workflows
-
-Trace Python script execution and optimize functions based on real-world usage patterns.
-
----
-
-## `codeflash optimize`
-
-Trace a Python script's execution and optimize functions based on real-world usage.
-
-```bash
-codeflash optimize <script.py> [script_args]
-```
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Basic trace and optimize
-    codeflash optimize app.py
-
-    # With script arguments
-    codeflash optimize process.py --input data.csv --output results.json
-
-    # Custom trace output file
-    codeflash optimize app.py --output custom_trace.trace
-
-    # With timeout (30 seconds)
-    codeflash optimize long_running_script.py --timeout 30
-
-    # Limit function trace count
-    codeflash optimize app.py --max-function-count 50
-
-    # Specify config file
-    codeflash optimize app.py --config-file-path pyproject.toml
-
-    # Local only (no PR)
-    codeflash optimize app.py --no-pr
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Basic trace and optimize
-    codeflash optimize app.py
-
-    # With script arguments
-    codeflash optimize process.py --input data.csv --output results.json
-
-    # Custom trace output file
-    codeflash optimize app.py --output custom_trace.trace
-
-    # With timeout (30 seconds)
-    codeflash optimize long_running_script.py --timeout 30
-
-    # Limit function trace count
-    codeflash optimize app.py --max-function-count 50
-
-    # Specify config file
-    codeflash optimize app.py --config-file-path pyproject.toml
-
-    # Local only (no PR)
-    codeflash optimize app.py --no-pr
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
-**How it works:**
-
-1. Runs your script with the provided arguments
-2. Traces all function calls during execution
-3. Identifies which functions are called and how often
-4. Generates replay tests based on actual usage
-5. Optimizes the traced functions
-
----
-
-## Trace with pytest
-
-Optimize functions called during pytest test execution.
-
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Trace pytest tests
-    codeflash optimize -m pytest tests/
-
-    # Trace specific test file
-    codeflash optimize -m pytest tests/test_core.py
-
-    # With pytest arguments
-    codeflash optimize -m pytest tests/ -v --tb=short
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Trace pytest tests
-    codeflash optimize -m pytest tests\
-
-    # Trace specific test file
-    codeflash optimize -m pytest tests\test_core.py
-
-    # With pytest arguments
-    codeflash optimize -m pytest tests\ -v --tb=short
-    ```
-  </Tab>
-</Tabs>
-
-<Tip>
-  Tracing pytest tests is great for optimizing functions that are heavily used in your test suite, ensuring optimizations work correctly with your existing tests.
-</Tip>
-
----
-
-## Trace Only (Generate Replay Tests)
-
-Create trace files and replay tests without running optimization.
-
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Trace only - generates replay test
-    codeflash optimize app.py --output trace_file.trace --trace-only
-
-    # Then optimize using the replay test
-    codeflash --replay-test tests/test_app_py__replay_test_0.py --no-pr
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Trace only - generates replay test
-    codeflash optimize app.py --output trace_file.trace --trace-only
-
-    # Then optimize using the replay test
-    codeflash --replay-test tests\test_app_py__replay_test_0.py --no-pr
-    ```
-  </Tab>
-</Tabs>
-
-<Note>
-  **Replay test naming**: Files are named based on the traced script path. For
-  `src/app.py`, the replay test will be named like
-  `test_srcapp_py__replay_test_0.py`.
-</Note>
-
-**Use cases for trace-only:**
-
-- Generate replay tests for later optimization
-- Debug tracing issues without running full optimization
-- Create reusable test cases from script execution
-
----
-
-## Replay Test Optimization
-
-Optimize functions using previously generated replay tests.
-
-```bash
-codeflash --replay-test <path/to/replay_test.py>
-```
-
-<Accordion title="Complete Examples">
-<Tabs>
-  <Tab title="Linux/macOS">
-    ```bash
-    # Optimize using replay test
-    codeflash --replay-test tests/test_app_py__replay_test_0.py --no-pr
-
-    # Multiple replay tests
-    codeflash --replay-test tests/test_*.py --no-pr
-    ```
-  </Tab>
-  <Tab title="Windows">
-    ```powershell
-    # Optimize using replay test
-    codeflash --replay-test tests\test_app_py__replay_test_0.py --no-pr
-
-    # Multiple replay tests (use Get-ChildItem for globbing)
-    codeflash --replay-test (Get-ChildItem tests\test_*.py) --no-pr
-    ```
-  </Tab>
-</Tabs>
-</Accordion>
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Optimization Commands"
-    icon="bullseye"
-    href="/cli-reference/optimization"
-  >
-    Learn about function optimization
-  </Card>
-  <Card
-    title="Flags Reference"
-    icon="list"
-    href="/cli-reference/flags"
-  >
-    Complete flag reference
-  </Card>
-</CardGroup>
-

docs/cli-reference/troubleshooting.mdx

@@ -1,157 +0,0 @@
----
-title: "CLI Troubleshooting"
-description: "Solutions for common Codeflash CLI issues and errors"
-icon: "wrench"
-sidebarTitle: "Troubleshooting"
-keywords: ["troubleshooting", "errors", "issues", "problems", "debugging"]
----
-
-# CLI Troubleshooting
-
-Solutions for common issues when using the Codeflash CLI.
-
----
-
-## Common Issues
-
-<AccordionGroup>
-  <Accordion title="'Functions outside module-root' Error">
-    **Problem**: Function not found because file is outside `module-root`.
-    
-    **Solution**: Ensure your file is within the `module-root` directory specified in `pyproject.toml`.
-    
-    ```bash
-    # Check your module-root
-    grep "module-root" pyproject.toml
-    
-    # Use the correct path (e.g., if module-root is "src")
-    codeflash --file src/myfile.py --function my_function --no-pr
-    ```
-  </Accordion>
-  
-  <Accordion title="'benchmarks-root must be specified' Error">
-    **Problem**: Using `--benchmark` without specifying benchmarks directory.
-    
-    **Solution**: Either add `benchmarks-root` to `pyproject.toml` or use the flag:
-    
-    ```bash
-    codeflash --file src/app.py --benchmark --benchmarks-root tests/benchmarks --no-pr
-    ```
-  </Accordion>
-  
-  <Accordion title="Replay Test File Not Found">
-    **Problem**: Replay test filename doesn't match expected path.
-    
-    **Solution**: Replay tests include the module path in their name. Check the actual filename:
-    
-    ```bash
-    # Linux/macOS
-    ls tests/test_*replay*.py
-    
-    # Windows
-    dir tests\test_*replay*.py
-    ```
-    
-    <Note>
-      Replay test files are named based on the traced script path. For `src/app.py`, 
-      the replay test will be named like `test_srcapp_py__replay_test_0.py`.
-    </Note>
-  </Accordion>
-  
-  <Accordion title="GitHub App Required">
-    **Problem**: PR creation fails due to missing GitHub App.
-    
-    **Solution**: Install the Codeflash GitHub App or use `--no-pr` for local optimization:
-    
-    ```bash
-    # Local optimization
-    codeflash --file src/app.py --function main --no-pr
-    
-    # Or install the GitHub App
-    # https://github.com/apps/codeflash-ai/installations/select_target
-    ```
-  </Accordion>
-
-  <Accordion title="Module Not Found Errors">
-    **Problem**: Codeflash can't find your Python modules.
-    
-    **Solution**: 
-    
-    1. Verify `module-root` is correctly set in `pyproject.toml`
-    2. Ensure you're running from the project root
-    3. Check that your Python environment has all dependencies installed
-    
-    ```bash
-    # Verify module-root
-    cat pyproject.toml | grep module-root
-    
-    # Check Python path
-    python -c "import sys; print(sys.path)"
-    ```
-  </Accordion>
-
-  <Accordion title="Test Generation Fails">
-    **Problem**: Codeflash can't generate tests for your function.
-    
-    **Solution**: 
-    
-    1. Ensure your function has a return statement
-    2. Check that the function is not a property or class method with special decorators
-    3. Use `--no-gen-tests` to skip test generation and use existing tests only
-    
-    ```bash
-    codeflash --file src/app.py --function main --no-gen-tests --no-pr
-    ```
-  </Accordion>
-
-  <Accordion title="Optimization Timeout">
-    **Problem**: Optimization takes too long or times out.
-    
-    **Solution**: 
-    
-    1. Use `--verbose` to see what's happening
-    2. For tracing, use `--timeout` to limit trace duration
-    3. For large functions, consider breaking them down
-    
-    ```bash
-    # Limit trace time
-    codeflash optimize app.py --timeout 30
-    
-    # See detailed progress
-    codeflash --file src/app.py --function main --verbose --no-pr
-    ```
-  </Accordion>
-</AccordionGroup>
-
----
-
-## Getting Help
-
-If you're still experiencing issues:
-
-1. **Check the logs**: Use `--verbose` flag to see detailed output
-2. **Verify setup**: Run `codeflash --verify-setup` to check your installation
-3. **Check configuration**: Ensure `pyproject.toml` is correctly configured
-4. **View help**: Run `codeflash --help` or `codeflash <command> --help`
-
----
-
-## Next Steps
-
-<CardGroup cols={2}>
-  <Card
-    title="Setup Commands"
-    icon="wrench"
-    href="/cli-reference/setup"
-  >
-    Review setup and initialization
-  </Card>
-  <Card
-    title="Flags Reference"
-    icon="list"
-    href="/cli-reference/flags"
-  >
-    Complete flag reference
-  </Card>
-</CardGroup>
-

docs/docs.json

@@ -18,35 +18,15 @@
       {
         "tab": "Documentation",
         "groups": [
-          {
-            "group": "🚀 Quickstart",
-            "pages": [
-              "getting-started/local-installation",
-              "getting-started/javascript-installation"
-            ]
-          },
           {
             "group": "🏠 Overview",
             "pages": ["index"]
           },
           {
-            "group": "📖 Codeflash CLI",
+            "group": "🚀 Quickstart",
             "pages": [
-              "cli-reference/index",
+              "getting-started/local-installation",
-              "cli-reference/setup",
+              "getting-started/javascript-installation"
-              "cli-reference/optimization",
-              "cli-reference/tracing",
-              "cli-reference/flags",
-              "cli-reference/troubleshooting"
-            ]
-          },
-          {
-            "group": "🛠 IDE Extension",
-            "pages": [
-              "editor-plugins/vscode/index",
-              "editor-plugins/vscode/features",
-              "editor-plugins/vscode/configuration",
-              "editor-plugins/vscode/troubleshooting"
             ]
           },
           {
@@ -65,6 +45,15 @@
               "optimizing-with-codeflash/review-optimizations"
             ]
           },
+          {
+            "group": "🛠 IDE Extension",
+            "pages": [
+              "editor-plugins/vscode/index",
+              "editor-plugins/vscode/features",
+              "editor-plugins/vscode/configuration",
+              "editor-plugins/vscode/troubleshooting"
+            ]
+          },
           {
             "group": "🧠 Core Concepts",
             "pages": [

docs/editor-plugins/vscode/configuration.mdx

@@ -146,7 +146,4 @@ When configuration issues are detected, the extension displays clear error messa
   <Card title="Project Configuration" icon="file-code" href="/configuration">
     Complete pyproject.toml reference
   </Card>
-  <Card title="Codeflash CLI" icon="terminal" href="/cli-reference/index">
-    Command-line options
-  </Card>
 </CardGroup>

docs/editor-plugins/vscode/index.mdx

@@ -204,7 +204,6 @@ The extension works alongside the Codeflash CLI. You can:
 - **Use extension for interactive work** — Optimize individual functions as you code
 - **Mix both** — The extension picks up CLI results when you return to the editor
 
-For CLI documentation, see the [Codeflash CLI](/cli-reference/index).
 
 ---
 
@@ -220,9 +219,6 @@ For CLI documentation, see the [Codeflash CLI](/cli-reference/index).
   <Card title="Troubleshooting" icon="wrench" href="/editor-plugins/vscode/troubleshooting">
     Fix common issues
   </Card>
-  <Card title="Codeflash CLI" icon="terminal" href="/cli-reference/index">
-    Command-line interface docs
-  </Card>
 </CardGroup>
 
 ---

docs/editor-plugins/vscode/troubleshooting.mdx

@@ -208,8 +208,5 @@ If you're still experiencing issues:
   <Card title="Configuration" icon="gear" href="/editor-plugins/vscode/configuration">
     Customize extension settings
   </Card>
-  <Card title="Codeflash CLI" icon="terminal" href="/cli-reference/index">
-    Command-line interface docs
-  </Card>
 </CardGroup>
 

tests/test_code_replacement.py

@@ -2119,7 +2119,6 @@ print("Hello world")
     expected_code = """import numpy as np
 
 a = 6
-
 if 2<3:
     a=4
 else:

tests/test_codeflash_capture.py

@@ -1602,7 +1602,94 @@ def calculate_portfolio_metrics(
         # now the test should match and no diffs should be found
         assert len(diffs) == 0
         assert matched
-        
+
     finally:
         test_path.unlink(missing_ok=True)
-        fto_file_path.unlink(missing_ok=True)
+        fto_file_path.unlink(missing_ok=True)
+
+
+def test_codeflash_capture_with_slots_class() -> None:
+    """Test that codeflash_capture works with classes that use __slots__ instead of __dict__."""
+    test_code = """
+from code_to_optimize.tests.pytest.sample_code import SlotsClass
+import unittest
+
+def test_slots_class():
+    obj = SlotsClass(10, "test")
+    assert obj.x == 10
+    assert obj.y == "test"
+"""
+    test_dir = (Path(__file__).parent.parent / "code_to_optimize" / "tests" / "pytest").resolve()
+    tmp_dir_path = get_run_tmp_file(Path("test_return_values"))
+    sample_code = f"""
+from codeflash.verification.codeflash_capture import codeflash_capture
+
+class SlotsClass:
+    __slots__ = ('x', 'y')
+
+    @codeflash_capture(function_name="SlotsClass.__init__", tmp_dir_path="{tmp_dir_path.as_posix()}", tests_root="{test_dir.as_posix()}")
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+"""
+    test_file_name = "test_slots_class_temp.py"
+    test_path = test_dir / test_file_name
+    test_path_perf = test_dir / "test_slots_class_temp_perf.py"
+
+    tests_root = Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/"
+    project_root_path = (Path(__file__).parent / "..").resolve()
+    sample_code_path = test_dir / "sample_code.py"
+
+    try:
+        with test_path.open("w") as f:
+            f.write(test_code)
+        with sample_code_path.open("w") as f:
+            f.write(sample_code)
+
+        test_env = os.environ.copy()
+        test_env["CODEFLASH_TEST_ITERATION"] = "0"
+        test_env["CODEFLASH_LOOP_INDEX"] = "1"
+        test_type = TestType.EXISTING_UNIT_TEST
+        test_config = TestConfig(
+            tests_root=tests_root,
+            tests_project_rootdir=project_root_path,
+            project_root_path=project_root_path,
+            test_framework="pytest",
+            pytest_cmd="pytest",
+        )
+        fto = FunctionToOptimize(
+            function_name="__init__",
+            file_path=sample_code_path,
+            parents=[FunctionParent(name="SlotsClass", type="ClassDef")],
+        )
+        func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
+        func_optimizer.test_files = TestFiles(
+            test_files=[
+                TestFile(
+                    instrumented_behavior_file_path=test_path,
+                    test_type=test_type,
+                    original_file_path=test_path,
+                    benchmarking_file_path=test_path_perf,
+                )
+            ]
+        )
+        test_results, coverage_data = func_optimizer.run_and_parse_tests(
+            testing_type=TestingMode.BEHAVIOR,
+            test_env=test_env,
+            test_files=func_optimizer.test_files,
+            optimization_iteration=0,
+            pytest_min_loops=1,
+            pytest_max_loops=1,
+            testing_time=0.1,
+        )
+
+        # Test should pass and capture the slots values
+        assert len(test_results) == 1
+        assert test_results[0].did_pass
+        # The return value should contain the slot values
+        assert test_results[0].return_value[0]["x"] == 10
+        assert test_results[0].return_value[0]["y"] == "test"
+
+    finally:
+        test_path.unlink(missing_ok=True)
+        sample_code_path.unlink(missing_ok=True)

tests/test_comparator.py

@@ -2316,6 +2316,77 @@ def test_dict_views() -> None:
     assert not comparator(d.items(), [("a", 1), ("b", 2)])
 
 
+def test_mappingproxy() -> None:
+    """Test comparator support for types.MappingProxyType (read-only dict view)."""
+    import types
+
+    # Basic equality
+    mp1 = types.MappingProxyType({"a": 1, "b": 2, "c": 3})
+    mp2 = types.MappingProxyType({"a": 1, "b": 2, "c": 3})
+    assert comparator(mp1, mp2)
+
+    # Different values
+    mp3 = types.MappingProxyType({"a": 1, "b": 2, "c": 4})
+    assert not comparator(mp1, mp3)
+
+    # Different keys
+    mp4 = types.MappingProxyType({"a": 1, "b": 2, "d": 3})
+    assert not comparator(mp1, mp4)
+
+    # Different length
+    mp5 = types.MappingProxyType({"a": 1, "b": 2})
+    assert not comparator(mp1, mp5)
+
+    # Order doesn't matter (like dict)
+    mp6 = types.MappingProxyType({"c": 3, "a": 1, "b": 2})
+    assert comparator(mp1, mp6)
+
+    # Empty mappingproxy
+    empty1 = types.MappingProxyType({})
+    empty2 = types.MappingProxyType({})
+    assert comparator(empty1, empty2)
+
+    # Nested values
+    nested1 = types.MappingProxyType({"a": [1, 2, 3], "b": {"x": 1}})
+    nested2 = types.MappingProxyType({"a": [1, 2, 3], "b": {"x": 1}})
+    nested3 = types.MappingProxyType({"a": [1, 2, 4], "b": {"x": 1}})
+    assert comparator(nested1, nested2)
+    assert not comparator(nested1, nested3)
+
+    # mappingproxy is not equal to dict (different types)
+    d = {"a": 1, "b": 2}
+    mp = types.MappingProxyType({"a": 1, "b": 2})
+    assert not comparator(mp, d)
+    assert not comparator(d, mp)
+
+    # Verify class __dict__ is indeed a mappingproxy
+    class MyClass:
+        x = 1
+        y = 2
+
+    assert isinstance(MyClass.__dict__, types.MappingProxyType)
+
+
+def test_mappingproxy_superset() -> None:
+    """Test comparator superset_obj support for mappingproxy."""
+    import types
+
+    mp1 = types.MappingProxyType({"a": 1, "b": 2})
+    mp2 = types.MappingProxyType({"a": 1, "b": 2, "c": 3})
+
+    # mp2 is a superset of mp1
+    assert comparator(mp1, mp2, superset_obj=True)
+    # mp1 is not a superset of mp2
+    assert not comparator(mp2, mp1, superset_obj=True)
+
+    # Same mappingproxy with superset_obj=True
+    assert comparator(mp1, mp1, superset_obj=True)
+
+    # Different values even with superset
+    mp3 = types.MappingProxyType({"a": 1, "b": 99, "c": 3})
+    assert not comparator(mp1, mp3, superset_obj=True)
+
+
 def test_tensorflow_tensor() -> None:
     """Test comparator support for TensorFlow Tensor objects."""
     try:

tests/test_get_read_writable_code.py

@@ -218,8 +218,28 @@ def test_no_targets_found() -> None:
             def target(self):
                 pass
     """
+    result = parse_code_and_prune_cst(dedent(code),CodeContextType.READ_WRITABLE, {"MyClass.Inner.target"})
+    expected = dedent("""
+    class MyClass:
+        def method(self):
+            pass
+
+        class Inner:
+            def target(self):
+                pass
+    """)
+    assert result.strip() == expected.strip()
+
+
+def test_no_targets_found_raises_for_nonexistent() -> None:
+    """Test that ValueError is raised when the target function doesn't exist at all."""
+    code = """
+    class MyClass:
+        def method(self):
+            pass
+    """
     with pytest.raises(ValueError, match="No target functions found in the provided code"):
-        parse_code_and_prune_cst(dedent(code),CodeContextType.READ_WRITABLE, {"MyClass.Inner.target"})
+        parse_code_and_prune_cst(dedent(code),CodeContextType.READ_WRITABLE, {"NonExistent.target"})
 
 
 def test_module_var() -> None:

tests/test_multi_file_code_replacement.py

@@ -124,7 +124,7 @@ def _get_string_usage(text: str) -> Usage:
     
     helper_file.unlink(missing_ok=True)
     main_file.unlink(missing_ok=True)
-    
+
     expected_helper = """import re
 from collections.abc import Sequence
 

tests/test_remove_unused_definitions.py

@@ -481,3 +481,86 @@ def unused_function():
     qualified_functions = {"get_platform_info", "get_loop_result"}
     result = remove_unused_definitions_by_function_names(code, qualified_functions)
     assert result.strip() == expected.strip()
+
+
+def test_enum_attribute_access_dependency() -> None:
+    """Test that enum/class attribute access like MessageKind.VALUE is tracked as a dependency."""
+    code = """
+from enum import Enum
+
+class MessageKind(Enum):
+    VALUE = "value"
+    OTHER = "other"
+
+class UnusedEnum(Enum):
+    UNUSED = "unused"
+
+UNUSED_VAR = 123
+
+def process_message(kind):
+    match kind:
+        case MessageKind.VALUE:
+            return "got value"
+        case MessageKind.OTHER:
+            return "got other"
+    return "unknown"
+"""
+
+    expected = """
+from enum import Enum
+
+class MessageKind(Enum):
+    VALUE = "value"
+    OTHER = "other"
+
+class UnusedEnum(Enum):
+    UNUSED = "unused"
+
+def process_message(kind):
+    match kind:
+        case MessageKind.VALUE:
+            return "got value"
+        case MessageKind.OTHER:
+            return "got other"
+    return "unknown"
+"""
+
+    qualified_functions = {"process_message"}
+    result = remove_unused_definitions_by_function_names(code, qualified_functions)
+    # MessageKind should be preserved because process_message uses MessageKind.VALUE
+    assert "class MessageKind" in result
+    # UNUSED_VAR should be removed
+    assert "UNUSED_VAR" not in result
+    assert result.strip() == expected.strip()
+
+
+def test_attribute_access_does_not_track_attr_name() -> None:
+    """Test that self.x attribute access doesn't track 'x' as a dependency on module-level x."""
+    code = """
+x = "module_level_x"
+UNUSED_VAR = "unused"
+
+class MyClass:
+    def __init__(self):
+        self.x = 1  # This 'x' is an attribute, not a reference to module-level 'x'
+
+    def get_x(self):
+        return self.x  # This 'x' is also an attribute access
+"""
+
+    expected = """
+class MyClass:
+    def __init__(self):
+        self.x = 1  # This 'x' is an attribute, not a reference to module-level 'x'
+
+    def get_x(self):
+        return self.x  # This 'x' is also an attribute access
+"""
+
+    qualified_functions = {"MyClass.get_x", "MyClass.__init__"}
+    result = remove_unused_definitions_by_function_names(code, qualified_functions)
+    # Module-level x should NOT be kept (self.x doesn't reference it)
+    assert 'x = "module_level_x"' not in result
+    # UNUSED_VAR should also be removed
+    assert "UNUSED_VAR" not in result
+    assert result.strip() == expected.strip()