⚡️ Speed up function normalize_code by 22% in PR #733 (deduplicate-better)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #733

If you approve this dependent PR, these changes will be merged into the original PR branch deduplicate-better.

This PR will be automatically closed if the original PR is merged.

---

📄 22% (0.22x) speedup for normalize_code in codeflash/code_utils/deduplicate_code.py

⏱️ Runtime : 96.6 milliseconds → 79.2 milliseconds (best of 72 runs)

📝 Explanation and details

The optimization replaces the remove_docstrings_from_ast function with fast_remove_docstrings_from_ast that uses a more efficient traversal strategy.

Key optimizations:

1. Eliminates ast.walk() overhead: The original code uses ast.walk() which visits every single node in the AST tree (21,611 hits in profiler). The optimized version uses a custom stack-based traversal that only visits nodes that can actually contain docstrings.

2. Targeted traversal: Instead of examining all AST nodes, the optimized version only traverses FunctionDef, AsyncFunctionDef, ClassDef, and Module nodes - the only node types that can contain docstrings in their body[0] position.

3. Reduced function call overhead: The stack-based approach eliminates the overhead of ast.walk()'s generator-based iteration, reducing the number of Python function calls from 21,611 to just the nodes that matter.

Performance impact: The docstring removal step drops from 131.4ms (25.5% of total time) to just 3.07ms (0.8% of total time) - a 97.7% reduction in that specific operation.

Test case effectiveness: The optimization shows consistent 10-25% speedups across all test cases, with the largest gains (23-24%) appearing in tests with many variables or docstrings (test_large_many_variables_*, test_large_docstring_removal_scaling). Even simple cases benefit from the reduced AST traversal overhead.

The optimization is particularly effective for code with deep nesting or many function/class definitions, as it avoids visiting irrelevant leaf nodes like literals, operators, and expressions that cannot contain docstrings.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 7 Passed
🌀 Generated Regression Tests	✅ 87 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_code_deduplication.py::test_deduplicate1	1.40ms	1.21ms	16.1%✅

🌀 Generated Regression Tests and Runtime

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.deduplicate_code import normalize_code


class VariableNormalizer(ast.NodeTransformer):
    """
    Normalizes variable names in assignments, but preserves function/class names and parameter names.
    Each scope (function, class, module) gets its own variable numbering.
    """
    def __init__(self):
        super().__init__()
        self.scope_stack = []
        self.global_var_counter = {}

    def visit_Module(self, node):
        self.scope_stack.append({})
        self.global_var_counter[id(node)] = 0
        self.generic_visit(node)
        self.scope_stack.pop()
        return node

    def visit_FunctionDef(self, node):
        self.scope_stack.append({})
        self.global_var_counter[id(node)] = 0
        # Don't normalize function name or args
        self.generic_visit(node)
        self.scope_stack.pop()
        return node

    def visit_AsyncFunctionDef(self, node):
        self.scope_stack.append({})
        self.global_var_counter[id(node)] = 0
        self.generic_visit(node)
        self.scope_stack.pop()
        return node

    def visit_ClassDef(self, node):
        self.scope_stack.append({})
        self.global_var_counter[id(node)] = 0
        self.generic_visit(node)
        self.scope_stack.pop()
        return node

    def visit_Name(self, node):
        # Only normalize variable names in assignment contexts (Store, Del, Load)
        if isinstance(node.ctx, (ast.Store, ast.Load, ast.Del)):
            # Do not normalize builtins, function names, class names, or parameters
            # Only normalize variable names in assignments and usage
            # Parameters are handled by ast.arg nodes, not ast.Name
            # So we can normalize all Name nodes except those that are function/class names
            # To avoid normalizing function/class names, check parent context (handled by visit_FunctionDef etc)
            # We'll normalize all variable names except those in function/class definitions

            # Find current scope
            scope = self.scope_stack[-1]
            counter = self.global_var_counter[list(self.global_var_counter.keys())[-1]]

            if node.id not in scope:
                # Assign next normalized name
                normalized_name = f"var_{len(scope)+1}"
                scope[node.id] = normalized_name
            node.id = scope[node.id]
        return node

    def visit_arg(self, node):
        # Don't normalize parameter names
        return node

    def visit_Attribute(self, node):
        # Don't normalize attribute names
        self.generic_visit(node)
        return node

    def visit_keyword(self, node):
        # Don't normalize keyword argument names
        self.generic_visit(node)
        return node


# unit tests

# ----------- BASIC TEST CASES -----------

def test_basic_assignment_normalization():
    # Single variable assignment
    code = "x = 1"
    expected = "var_1 = 1"
    codeflash_output = normalize_code(code) # 73.4μs -> 66.2μs (10.8% faster)

def test_basic_multiple_assignments():
    # Multiple assignments in sequence
    code = "x = 1\ny = 2\nz = x + y"
    expected = "var_1 = 1\nvar_2 = 2\nvar_3 = var_1 + var_2"
    codeflash_output = normalize_code(code) # 128μs -> 111μs (15.1% faster)

def test_basic_function_preserves_names_and_params():
    # Function name and parameter should be preserved, variable inside normalized
    code = "def foo(a):\n    b = a + 1\n    return b"
    expected = "def foo(a):\n    var_1 = a + 1\n    return var_1"
    codeflash_output = normalize_code(code) # 146μs -> 130μs (11.9% faster)

def test_basic_class_preserves_names_and_attrs():
    # Class name and attribute names should be preserved
    code = "class MyClass:\n    def __init__(self, x):\n        self.x = x"
    expected = "class MyClass:\n    def __init__(self, x):\n        self.x = x"
    codeflash_output = normalize_code(code) # 136μs -> 123μs (10.6% faster)

def test_basic_docstring_removal():
    # Docstring should be removed by default
    code = '"""Module docstring"""\nx = 1'
    expected = "var_1 = 1"
    codeflash_output = normalize_code(code) # 74.5μs -> 65.5μs (13.8% faster)

def test_basic_docstring_preserved_if_requested():
    # Docstring should be preserved if remove_docstrings=False
    code = '"""Module docstring"""\nx = 1'
    expected = '"""Module docstring"""\nvar_1 = 1'
    codeflash_output = normalize_code(code, remove_docstrings=False) # 76.5μs -> 76.4μs (0.064% faster)

def test_basic_function_docstring_removal():
    # Function docstring should be removed
    code = 'def foo():\n    """Docstring"""\n    x = 1\n    return x'
    expected = 'def foo():\n    var_1 = 1\n    return var_1'
    codeflash_output = normalize_code(code) # 122μs -> 107μs (14.1% faster)

def test_basic_function_docstring_preserved():
    # Function docstring should be preserved if requested
    code = 'def foo():\n    """Docstring"""\n    x = 1\n    return x'
    expected = 'def foo():\n    """Docstring"""\n    var_1 = 1\n    return var_1'
    codeflash_output = normalize_code(code, remove_docstrings=False) # 117μs -> 116μs (1.35% faster)

def test_basic_attribute_access_not_normalized():
    # Attribute names should not be normalized
    code = "class A:\n    def foo(self):\n        self.x = 1\n        return self.x"
    expected = "class A:\n    def foo(self):\n        self.x = 1\n        return self.x"
    codeflash_output = normalize_code(code) # 155μs -> 138μs (12.8% faster)

def test_basic_keyword_args_not_normalized():
    # Keyword argument names should not be normalized
    code = "def foo(x=1, y=2):\n    return x + y"
    expected = "def foo(x=1, y=2):\n    return x + y"
    codeflash_output = normalize_code(code) # 141μs -> 125μs (12.8% faster)

# ----------- EDGE TEST CASES -----------

def test_edge_empty_code():
    # Empty string should return empty string
    code = ""
    expected = ""
    codeflash_output = normalize_code(code) # 33.4μs -> 29.9μs (11.7% faster)

def test_edge_invalid_syntax():
    # Invalid syntax should raise ValueError
    code = "def foo("
    with pytest.raises(ValueError):
        normalize_code(code) # 21.4μs -> 21.3μs (0.606% faster)

def test_edge_single_variable():
    # Single variable, single line
    code = "x=1"
    expected = "var_1 = 1"
    codeflash_output = normalize_code(code) # 69.3μs -> 62.3μs (11.3% faster)

def test_edge_variable_shadowing_in_functions():
    # Variables with same name in different scopes should be normalized independently
    code = "x = 1\ndef foo():\n    x = 2\n    return x\nz = x"
    expected = "var_1 = 1\ndef foo():\n    var_1 = 2\n    return var_1\nvar_2 = var_1"
    codeflash_output = normalize_code(code) # 157μs -> 137μs (15.0% faster)

def test_edge_nested_functions():
    # Nested function should have independent variable normalization
    code = "def outer():\n    x = 1\n    def inner():\n        x = 2\n        return x\n    return x"
    expected = "def outer():\n    var_1 = 1\n    def inner():\n        var_1 = 2\n        return var_1\n    return var_1"
    codeflash_output = normalize_code(code) # 170μs -> 150μs (13.7% faster)

def test_edge_multiple_assignments_same_line():
    # Multiple assignments in one line
    code = "x = y = 1"
    expected = "var_1 = var_2 = 1"
    codeflash_output = normalize_code(code) # 73.1μs -> 65.2μs (12.1% faster)

def test_edge_tuple_unpacking():
    # Tuple unpacking assignment
    code = "a, b = 1, 2"
    expected = "var_1, var_2 = 1, 2"
    codeflash_output = normalize_code(code) # 102μs -> 91.3μs (12.7% faster)

def test_edge_list_unpacking():
    # List unpacking assignment
    code = "[a, b] = [1, 2]"
    expected = "[var_1, var_2] = [1, 2]"
    codeflash_output = normalize_code(code) # 103μs -> 92.3μs (12.7% faster)

def test_edge_for_loop_variable_normalized():
    # For loop variable should be normalized
    code = "for x in range(5):\n    print(x)"
    expected = "for var_1 in range(5):\n    print(var_1)"
    codeflash_output = normalize_code(code) # 118μs -> 106μs (11.8% faster)

def test_edge_with_statement_variable_normalized():
    # With statement variable should be normalized
    code = "with open('f') as f:\n    data = f.read()"
    expected = "with open('f') as var_1:\n    var_2 = var_1.read()"
    codeflash_output = normalize_code(code) # 134μs -> 119μs (12.5% faster)

def test_edge_del_statement_normalized():
    # del statement variable should be normalized
    code = "x = 1\ndel x"
    expected = "var_1 = 1\ndel var_1"
    codeflash_output = normalize_code(code) # 80.4μs -> 71.2μs (12.8% faster)

def test_edge_augmented_assignment_normalized():
    # Augmented assignment variable should be normalized
    code = "x = 1\nx += 2"
    expected = "var_1 = 1\nvar_1 += 2"
    codeflash_output = normalize_code(code) # 88.2μs -> 79.3μs (11.2% faster)

def test_edge_global_and_nonlocal_not_normalized():
    # global/nonlocal statements should not be normalized
    code = "def foo():\n    global x\n    x = 1"
    expected = "def foo():\n    global x\n    var_1 = 1"
    codeflash_output = normalize_code(code) # 109μs -> 98.8μs (10.7% faster)

def test_edge_lambda_preserves_param_names():
    # Lambda parameter names should be preserved, usage normalized
    code = "f = lambda x: x + 1"
    expected = "var_1 = lambda x: x + 1"
    codeflash_output = normalize_code(code) # 120μs -> 106μs (12.9% faster)

def test_edge_comprehension_variable_normalized():
    # List comprehension variable normalized
    code = "lst = [x for x in range(5)]"
    expected = "var_1 = [var_2 for var_2 in range(5)]"
    codeflash_output = normalize_code(code) # 120μs -> 107μs (12.8% faster)

def test_edge_set_and_dict_comprehension_variable_normalized():
    # Set and dict comprehension variable normalized
    code = "s = {x for x in range(3)}\nd = {x: x for x in range(3)}"
    expected = "var_1 = {var_2 for var_2 in range(3)}\nvar_3 = {var_4: var_4 for var_4 in range(3)}"
    codeflash_output = normalize_code(code) # 186μs -> 163μs (13.8% faster)

def test_edge_nested_list_comprehension():
    # Nested list comprehension variable normalization
    code = "lst = [[y for y in range(2)] for x in range(3)]"
    expected = "var_1 = [[var_3 for var_3 in range(2)] for var_2 in range(3)]"
    codeflash_output = normalize_code(code) # 159μs -> 138μs (15.6% faster)

def test_edge_try_except_variable_normalized():
    # Exception variable normalized
    code = "try:\n    x = 1\nexcept Exception as e:\n    print(e)"
    expected = "try:\n    var_1 = 1\nexcept Exception as var_2:\n    print(var_2)"
    codeflash_output = normalize_code(code) # 130μs -> 115μs (12.7% faster)

def test_edge_multiple_functions_and_classes():
    # Multiple functions and classes, each scope normalized independently
    code = "def foo():\n    x = 1\nclass Bar:\n    y = 2"
    expected = "def foo():\n    var_1 = 1\nclass Bar:\n    var_1 = 2"
    codeflash_output = normalize_code(code) # 141μs -> 126μs (11.4% faster)

def test_edge_async_function_normalization():
    # Async function normalization
    code = "async def foo():\n    x = await bar()\n    return x"
    expected = "async def foo():\n    var_1 = await bar()\n    return var_1"
    codeflash_output = normalize_code(code) # 130μs -> 116μs (12.4% faster)

def test_edge_function_with_kwargs_and_args():
    # Function with *args and **kwargs, parameter names preserved
    code = "def f(*args, **kwargs):\n    x = args[0]\n    y = kwargs['a']"
    expected = "def f(*args, **kwargs):\n    var_1 = args[0]\n    var_2 = kwargs['a']"
    codeflash_output = normalize_code(code) # 167μs -> 147μs (13.4% faster)

def test_edge_function_with_default_and_annotation():
    # Function with default and annotation, parameter names preserved
    code = "def f(x: int = 1):\n    y = x + 1"
    expected = "def f(x: int = 1):\n    var_1 = x + 1"
    codeflash_output = normalize_code(code) # 143μs -> 129μs (11.2% faster)

def test_edge_class_with_inheritance_and_method():
    # Class with inheritance and method, only method variable normalized
    code = "class A(B):\n    def foo(self):\n        x = 1\n        return x"
    expected = "class A(B):\n    def foo(self):\n        var_1 = 1\n        return var_1"
    codeflash_output = normalize_code(code) # 148μs -> 131μs (12.7% faster)

def test_edge_ann_assign_normalized():
    # Annotated assignment normalized
    code = "x: int = 1"
    expected = "var_1: int = 1"
    codeflash_output = normalize_code(code) # 72.7μs -> 64.6μs (12.5% faster)

def test_edge_multiline_string_assignment():
    # Multiline string assignment normalized
    code = "x = '''hello\nworld'''\ny = x"
    expected = "var_1 = '''hello\nworld'''\nvar_2 = var_1"
    codeflash_output = normalize_code(code) # 88.1μs -> 76.5μs (15.1% faster)

def test_edge_variable_in_if_else_blocks():
    # Variable in if/else blocks normalized
    code = "if True:\n    x = 1\nelse:\n    x = 2"
    expected = "if True:\n    var_1 = 1\nelse:\n    var_1 = 2"
    codeflash_output = normalize_code(code) # 111μs -> 99.4μs (12.0% faster)

def test_edge_variable_in_while_loop():
    # Variable in while loop normalized
    code = "while True:\n    x = 1"
    expected = "while True:\n    var_1 = 1"
    codeflash_output = normalize_code(code) # 86.7μs -> 77.4μs (12.0% faster)

def test_edge_variable_in_try_finally():
    # Variable in try/finally normalized
    code = "try:\n    x = 1\nfinally:\n    x = 2"
    expected = "try:\n    var_1 = 1\nfinally:\n    var_1 = 2"
    codeflash_output = normalize_code(code) # 104μs -> 92.5μs (12.9% faster)

def test_edge_variable_in_with_multiple_targets():
    # With statement with multiple targets
    code = "with a as x, b as y:\n    z = x + y"
    expected = "with var_1 as var_2, var_3 as var_4:\n    var_5 = var_2 + var_4"
    codeflash_output = normalize_code(code) # 128μs -> 111μs (15.9% faster)

def test_edge_variable_in_match_case():
    # Python 3.10+ match/case statement
    code = "match x:\n    case 1:\n        y = 2"
    expected = "match var_1:\n    case 1:\n        var_2 = 2"
    codeflash_output = normalize_code(code) # 107μs -> 94.9μs (13.7% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_large_many_variables_module_scope():
    # Large number of variables in module scope
    code_lines = [f"x{i} = {i}" for i in range(1, 501)]
    code = "\n".join(code_lines)
    expected_lines = [f"var_{i} = {i}" for i in range(1, 501)]
    expected = "\n".join(expected_lines)
    codeflash_output = normalize_code(code) # 8.54ms -> 6.94ms (23.1% faster)

def test_large_many_variables_function_scope():
    # Large number of variables in function scope
    code_lines = [f"x{i} = {i}" for i in range(1, 501)]
    code = "def foo():\n    " + "\n    ".join(code_lines)
    expected_lines = [f"var_{i} = {i}" for i in range(1, 501)]
    expected = "def foo():\n    " + "\n    ".join(expected_lines)
    codeflash_output = normalize_code(code) # 8.56ms -> 6.98ms (22.6% faster)

def test_large_many_variables_class_scope():
    # Large number of variables in class scope
    code_lines = [f"x{i} = {i}" for i in range(1, 501)]
    code = "class A:\n    " + "\n    ".join(code_lines)
    expected_lines = [f"var_{i} = {i}" for i in range(1, 501)]
    expected = "class A:\n    " + "\n    ".join(expected_lines)
    codeflash_output = normalize_code(code) # 8.47ms -> 6.96ms (21.7% faster)

def test_large_many_functions_and_classes():
    # Many functions and classes, each with a few variables
    code = ""
    expected = ""
    for i in range(1, 51):
        code += f"def f{i}():\n    x = {i}\n    y = x + 1\n"
        expected += f"def f{i}():\n    var_1 = {i}\n    var_2 = var_1 + 1\n"
    for i in range(1, 51):
        code += f"class C{i}:\n    z = {i}\n"
        expected += f"class C{i}:\n    var_1 = {i}\n"
    codeflash_output = normalize_code(code.strip()) # 4.99ms -> 4.16ms (19.9% faster)

def test_large_comprehension_scaling():
    # Large list comprehension
    code = "lst = [x for x in range(1000)]"
    expected = "var_1 = [var_2 for var_2 in range(1000)]"
    codeflash_output = normalize_code(code) # 127μs -> 113μs (12.1% faster)

def test_large_nested_comprehension_scaling():
    # Large nested list comprehension
    code = "lst = [[y for y in range(10)] for x in range(100)]"
    expected = "var_1 = [[var_3 for var_3 in range(10)] for var_2 in range(100)]"
    codeflash_output = normalize_code(code) # 163μs -> 142μs (14.8% faster)

def test_large_docstring_removal_scaling():
    # Large module with many docstrings
    code = '"""Docstring"""\n' + "\n".join([f'def f{i}():\n    """Docstring"""\n    x = {i}' for i in range(1, 101)])
    expected = "\n".join([f'def f{i}():\n    var_1 = {i}' for i in range(1, 101)])
    codeflash_output = normalize_code(code) # 4.30ms -> 3.49ms (23.2% faster)

def test_large_assignment_and_usage():
    # Large number of variables, each used in expressions
    code_lines = [f"x{i} = {i}" for i in range(1, 501)]
    code_lines += [f"y{i} = x{i} + 1" for i in range(1, 501)]
    code = "\n".join(code_lines)
    expected_lines = [f"var_{i} = {i}" for i in range(1, 501)]
    expected_lines += [f"var_{i+500} = var_{i} + 1" for i in range(1, 501)]
    expected = "\n".join(expected_lines)
    codeflash_output = normalize_code(code) # 23.5ms -> 18.9ms (24.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.deduplicate_code import normalize_code


class VariableNormalizer(ast.NodeTransformer):
    """
    Normalizes variable names in Python code.
    Only variable names are normalized; function/class names and parameters are preserved.
    """
    def __init__(self):
        super().__init__()
        self.var_counter = 0
        self.var_map_stack = [{}]
        self.global_vars = set()

    def _new_var(self):
        self.var_counter += 1
        return f"var_{self.var_counter}"

    def visit_FunctionDef(self, node):
        # New scope for local variables
        self.var_map_stack.append({})
        self.generic_visit(node)
        self.var_map_stack.pop()
        return node

    def visit_AsyncFunctionDef(self, node):
        self.var_map_stack.append({})
        self.generic_visit(node)
        self.var_map_stack.pop()
        return node

    def visit_ClassDef(self, node):
        # New scope for class variables
        self.var_map_stack.append({})
        self.generic_visit(node)
        self.var_map_stack.pop()
        return node

    def visit_Name(self, node):
        # Only normalize variable names (not function/class names, not attributes, not parameters)
        # Only normalize if it's not a builtin or keyword
        # Don't normalize if it's in global_vars
        if isinstance(node.ctx, (ast.Store, ast.Load, ast.Del)):
            # Check if it's a parameter or attribute
            if not self._is_param_or_attr(node):
                # Check if it's global
                if node.id not in self.global_vars:
                    var_map = self.var_map_stack[-1]
                    if node.id not in var_map:
                        var_map[node.id] = self._new_var()
                    node.id = var_map[node.id]
        return node

    def visit_Global(self, node):
        # Track global variables so we don't normalize them
        for name in node.names:
            self.global_vars.add(name)
        return node

    def visit_arg(self, node):
        # Don't normalize parameter names
        return node

    def visit_Attribute(self, node):
        # Don't normalize attribute names
        self.generic_visit(node)
        return node

    def _is_param_or_attr(self, node):
        # Check if node is a parameter or attribute
        # Parameters are handled in visit_arg, attributes in visit_Attribute
        return False

# unit tests

# --------------------------
# BASIC TEST CASES
# --------------------------

def test_basic_assignment():
    # Test that variable names are normalized in simple assignments
    code = "x = 1\ny = x + 2"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 112μs -> 99.1μs (13.3% faster)

def test_basic_function():
    # Test that variable names inside a function are normalized, but parameter and function name are preserved
    code = "def foo(a):\n    x = a + 1\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 145μs -> 129μs (12.1% faster)

def test_basic_class():
    # Test that variable names inside a class method are normalized, but class and method names are preserved
    code = "class Bar:\n    def baz(self):\n        x = 5\n        return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 142μs -> 125μs (13.5% faster)

def test_multiple_variables():
    # Test that multiple variables are normalized with unique names
    code = "a = 1\nb = 2\nc = a + b"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 124μs -> 109μs (14.2% faster)

def test_docstring_removal():
    # Test that docstrings are removed by default
    code = '"""Module docstring"""\ndef foo():\n    """Function docstring"""\n    x = 1\n    return x'
    codeflash_output = normalize_code(code); normalized = codeflash_output # 127μs -> 110μs (15.3% faster)

def test_docstring_preserved():
    # Test that docstrings are preserved when remove_docstrings=False
    code = '"""Module docstring"""\ndef foo():\n    """Function docstring"""\n    x = 1\n    return x'
    codeflash_output = normalize_code(code, remove_docstrings=False); normalized = codeflash_output # 132μs -> 131μs (0.913% faster)

def test_nested_functions():
    # Test that variable names are normalized independently in nested functions
    code = "def outer():\n    x = 1\n    def inner():\n        y = 2\n        return y\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 172μs -> 151μs (14.1% faster)

def test_global_variable():
    # Test that global variables are not normalized
    code = "x = 1\nglobal x\ndef foo():\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 123μs -> 108μs (13.9% faster)

def test_attribute_not_normalized():
    # Test that attribute names are not normalized
    code = "class A:\n    def f(self):\n        self.x = 5\n        return self.x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 154μs -> 135μs (13.6% faster)

def test_parameter_not_normalized():
    # Test that parameter names are preserved
    code = "def foo(bar):\n    x = bar + 1\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 143μs -> 126μs (13.0% faster)

def test_async_function():
    # Test that async function variables are normalized
    code = "async def foo():\n    x = 5\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 115μs -> 101μs (13.0% faster)

# --------------------------
# EDGE TEST CASES
# --------------------------

def test_empty_code():
    # Test that empty code returns empty string
    code = ""
    codeflash_output = normalize_code(code); normalized = codeflash_output # 33.4μs -> 29.6μs (12.8% faster)

def test_only_docstring():
    # Test that code with only a docstring is removed
    code = '"""Just a docstring"""'
    codeflash_output = normalize_code(code); normalized = codeflash_output # 41.3μs -> 35.1μs (17.8% faster)

def test_syntax_error():
    # Test that invalid syntax raises ValueError
    code = "def foo("
    with pytest.raises(ValueError):
        normalize_code(code) # 21.6μs -> 21.8μs (0.780% slower)

def test_variable_shadowing():
    # Test that variable shadowing in inner scopes is handled correctly
    code = "x = 1\ndef foo():\n    x = 2\n    return x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 143μs -> 124μs (15.4% faster)

def test_multiple_assignments():
    # Test multiple assignments on one line
    code = "a, b = 1, 2"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 104μs -> 92.0μs (13.9% faster)

def test_tuple_unpacking():
    # Test tuple unpacking in assignment
    code = "a, (b, c) = 1, (2, 3)"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 135μs -> 118μs (14.5% faster)

def test_list_comprehension():
    # Test variable normalization in list comprehensions
    code = "lst = [x for x in range(5)]"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 120μs -> 105μs (13.5% faster)

def test_lambda_variable():
    # Test variable normalization in lambda expressions
    code = "f = lambda x: x + 1"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 120μs -> 105μs (14.8% faster)

def test_nested_class():
    # Test variable normalization in nested classes
    code = "class A:\n    x = 1\n    class B:\n        y = 2"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 121μs -> 108μs (11.8% faster)

def test_del_statement():
    # Test normalization in del statements
    code = "x = 1\ndel x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 79.8μs -> 70.8μs (12.8% faster)

def test_augmented_assignment():
    # Test normalization in augmented assignment
    code = "x = 1\nx += 2"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 88.3μs -> 78.1μs (13.0% faster)

def test_for_loop():
    # Test normalization in for loop variable
    code = "for i in range(3):\n    x = i"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 111μs -> 97.2μs (15.0% faster)

def test_with_statement():
    # Test normalization in with statement variable
    code = "with open('f') as f:\n    x = f.read()"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 134μs -> 118μs (14.0% faster)

def test_try_except():
    # Test normalization in try/except block
    code = "try:\n    x = 1\nexcept Exception as e:\n    y = e"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 119μs -> 103μs (14.7% faster)

def test_nonlocal_variable():
    # Test nonlocal variable is not normalized
    code = "def outer():\n    x = 1\n    def inner():\n        nonlocal x\n        x = 2"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 158μs -> 139μs (13.3% faster)

def test_variable_in_generator_expression():
    # Test normalization in generator expressions
    code = "gen = (x for x in range(5))"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 119μs -> 104μs (14.2% faster)

def test_variable_in_set_dict_comprehension():
    # Test normalization in set and dict comprehensions
    code = "s = {x for x in range(5)}\nd = {x: x for x in range(5)}"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 185μs -> 160μs (15.7% faster)

def test_variable_in_while_loop():
    # Test normalization in while loop variable
    code = "x = 0\nwhile x < 5:\n    x += 1"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 132μs -> 117μs (12.3% faster)

def test_variable_in_if_statement():
    # Test normalization in if statement
    code = "x = 1\nif x > 0:\n    y = x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 127μs -> 110μs (15.0% faster)

def test_function_with_default_args():
    # Test that default argument expressions are not normalized
    code = "def foo(x=1):\n    y = x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 117μs -> 103μs (13.4% faster)

def test_function_with_kwargs():
    # Test that kwargs are preserved
    code = "def foo(**kwargs):\n    x = kwargs['a']"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 127μs -> 111μs (14.1% faster)

def test_function_with_args():
    # Test that *args are preserved
    code = "def foo(*args):\n    x = args[0]"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 124μs -> 109μs (13.2% faster)

def test_multiline_string():
    # Test that multiline strings are not treated as docstrings unless at top of module/class/function
    code = "x = '''multi\nline\nstring'''\ny = x"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 87.1μs -> 74.8μs (16.4% faster)

# --------------------------
# LARGE SCALE TEST CASES
# --------------------------

def test_large_number_of_variables():
    # Test normalization with a large number of variables
    code_lines = [f"x{i} = {i}" for i in range(1, 501)]
    code = "\n".join(code_lines)
    codeflash_output = normalize_code(code); normalized = codeflash_output # 8.54ms -> 6.88ms (24.2% faster)
    expected_lines = [f"var_{i} = {i}" for i in range(1, 501)]
    expected = "\n".join(expected_lines)

def test_large_nested_functions():
    # Test normalization with deeply nested functions and variables
    code = "def f0():\n"
    for i in range(1, 51):
        code += "    " * i + f"def f{i}():\n"
    code += "    " * 51 + "x = 1\n"
    code += "    " * 51 + "return x\n"
    for i in reversed(range(1, 52)):
        code += "    " * i + "return f{}\n".format(i if i < 51 else 51)
    codeflash_output = normalize_code(code); normalized = codeflash_output # 2.05ms -> 1.77ms (15.7% faster)

def test_large_comprehension():
    # Test normalization with a large list comprehension
    code = "lst = [x for x in range(1000)]"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 126μs -> 112μs (12.4% faster)

def test_large_class():
    # Test normalization with a large class containing many variables
    code = "class Big:\n"
    for i in range(1, 501):
        code += f"    x{i} = {i}\n"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 8.58ms -> 6.91ms (24.2% faster)
    expected = "class Big:\n" + "".join([f"    var_{i} = {i}\n" for i in range(1, 501)])

def test_large_function_body():
    # Test normalization with a large function body
    code = "def foo():\n"
    for i in range(1, 501):
        code += f"    x{i} = {i}\n"
    code += "    return x500"
    codeflash_output = normalize_code(code); normalized = codeflash_output # 8.62ms -> 7.01ms (23.0% faster)
    expected = "def foo():\n" + "".join([f"    var_{i} = {i}\n" for i in range(1, 501)]) + "    return var_500"
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from codeflash.code_utils.deduplicate_code import normalize_code
import pytest

def test_normalize_code():
    with pytest.raises(TypeError, match='compile\\(\\)\\ arg\\ 1\\ must\\ be\\ a\\ string,\\ bytes\\ or\\ AST\\ object'):
        normalize_code('', remove_docstrings=False)

To edit these changes git checkout codeflash/optimize-pr733-2025-09-13T23.50.11 and push.

[misrasaurabh1]

adapted this into the branch