fix overlappings args in codeflash wrap

[aseembits93]

PR Type

Bug fix, Tests

---

Description

• Prefix wrapper args to avoid overlaps

• Update AST and f-strings variable names

• Adjust DB insert fields to prefixed args

• Update tests to new wrapper signature

---

Diagram Walkthrough

flowchart LR
  wrap["codeflash_wrap signature"]
  astbuild["AST string/tag construction"]
  call["Wrapped function call"]
  db["DB insert payload"]
  tests["Tests updated strings"]

  wrap -- "prefixed arg names" --> astbuild
  wrap -- "prefixed arg names" --> call
  wrap -- "prefixed arg names" --> db
  wrap -- "new expectations" --> tests

Loading

File Walkthrough

	Relevant files
Bug fix	instrument_existing_tests.py AST wrapper refactor to prefixed argument names                    codeflash/code_utils/instrument_existing_tests.py Prefix all wrapper parameters with codeflash_. Update AST f-strings to use prefixed names. Call codeflash_wrapped instead of wrapped. Insert prefixed fields into test_results. +31/-25
Tests	test_instrument_all_and_run.py Align test wrapper string to prefixed args                              tests/test_instrument_all_and_run.py Update inline codeflash_wrap string signature. Replace f-string variables with prefixed versions. Use codeflash_wrapped in wrapped call. Update DB insert tuple to prefixed fields. +6/-6      test_instrument_tests.py Update instrumentation tests for new wrapper API                  tests/test_instrument_tests.py Update expected wrapper strings to prefixed args. Modify stdout tag construction variables. Use codeflash_wrapped for function invocation. Adjust DB insert fields to prefixed names. +23/-23

---

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Consistency Ensure all occurrences of the old argument and variable names (e.g., wrapped, function_name, line_id, loop_index) are fully replaced with their codeflash_-prefixed counterparts across the entire AST construction, especially in any branches not shown here (e.g., different TestingMode paths) to avoid NameError at runtime. args=[], keywords=[], ), lineno=lineno + 11, ), ast.Assign( targets=[ast.Name(id="return_value", ctx=ast.Store())], value=ast.Call( func=ast.Name(id="codeflash_wrapped", ctx=ast.Load()), args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())], keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))], ), Test Fixture Strings The embedded wrapper strings now use codeflash_-prefixed names; verify that all consumers of these strings in the tests and any CLI instrumentation paths expect the new signature and tag format, including perf-only variants, to prevent brittle string comparisons. codeflash_wrap_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, codeflash_cur, codeflash_con, *args, **kwargs): test_id = f'{{codeflash_test_module_name}}:{{codeflash_test_class_name}}:{{codeflash_test_name}}:{{codeflash_line_id}}:{{codeflash_loop_index}}' if not hasattr(codeflash_wrap, 'index'): codeflash_wrap.index = {{}} if test_id in codeflash_wrap.index: codeflash_wrap.index[test_id] += 1 else: codeflash_wrap.index[test_id] = 0 codeflash_test_index = codeflash_wrap.index[test_id] invocation_id = f'{{codeflash_line_id}}_{{codeflash_test_index}}' test_stdout_tag = f"{{codeflash_test_module_name}}:{{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}}{{codeflash_test_name}}:{{codeflash_function_name}}:{{codeflash_loop_index}}:{{invocation_id}}" print(f"!$######{{test_stdout_tag}}######$!") exception = None gc.disable() try: counter = time.perf_counter_ns() return_value = codeflash_wrapped(*args, **kwargs) codeflash_duration = time.perf_counter_ns() - counter except Exception as e: codeflash_duration = time.perf_counter_ns() - counter exception = e gc.enable() print(f"!######{{test_stdout_tag}}######!") pickled_return_value = pickle.dumps(exception) if exception else pickle.dumps(return_value) codeflash_cur.execute('INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)', (codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_loop_index, invocation_id, codeflash_duration, pickled_return_value, 'function_call')) codeflash_con.commit() if exception: raise exception return return_value """ codeflash_wrap_perfonly_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs): test_id = f'{{codeflash_test_module_name}}:{{codeflash_test_class_name}}:{{codeflash_test_name}}:{{codeflash_line_id}}:{{codeflash_loop_index}}' if not hasattr(codeflash_wrap, 'index'): codeflash_wrap.index = {{}} if test_id in codeflash_wrap.index: codeflash_wrap.index[test_id] += 1 else: codeflash_wrap.index[test_id] = 0 codeflash_test_index = codeflash_wrap.index[test_id] invocation_id = f'{{codeflash_line_id}}_{{codeflash_test_index}}' test_stdout_tag = f"{{codeflash_test_module_name}}:{{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}}{{codeflash_test_name}}:{{codeflash_function_name}}:{{codeflash_loop_index}}:{{invocation_id}}" print(f"!$######{{test_stdout_tag}}######$!") exception = None gc.disable() try: counter = time.perf_counter_ns() return_value = codeflash_wrapped(*args, **kwargs) codeflash_duration = time.perf_counter_ns() - counter except Exception as e: codeflash_duration = time.perf_counter_ns() - counter exception = e DB Schema Alignment The INSERT payload order changed to use prefixed identifiers; confirm this ordering matches the DB schema and any downstream readers that parse test_results, particularly for function_name and index/ids, to avoid data misalignment. value=ast.Call( func=ast.Attribute( value=ast.Name(id="codeflash_cur", ctx=ast.Load()), attr="execute", ctx=ast.Load() ), args=[ ast.Constant(value="INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"), ast.Tuple( elts=[ ast.Name(id="codeflash_test_module_name", ctx=ast.Load()), ast.Name(id="codeflash_test_class_name", ctx=ast.Load()), ast.Name(id="codeflash_test_name", ctx=ast.Load()), ast.Name(id="codeflash_function_name", ctx=ast.Load()), ast.Name(id="codeflash_loop_index", ctx=ast.Load()), ast.Name(id="invocation_id", ctx=ast.Load()), ast.Name(id="codeflash_duration", ctx=ast.Load()), ast.Name(id="pickled_return_value", ctx=ast.Load()), ast.Constant(value=VerificationType.FUNCTION_CALL.value),

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Re-enable GC reliably Ensure gc.enable() runs even if an unexpected error occurs before the try block exits. Move gc.enable() into a finally block to guarantee GC is re-enabled on all paths. tests/test_instrument_tests.py [74-82] try: counter = time.perf_counter_ns() return_value = codeflash_wrapped(*args, **kwargs) codeflash_duration = time.perf_counter_ns() - counter except Exception as e: codeflash_duration = time.perf_counter_ns() - counter exception = e -gc.enable() +finally: + gc.enable() Suggestion importance[1-10]: 7 __ Why: Moving gc.enable() to a finally block ensures GC is always re-enabled, improving reliability without changing behavior; it’s a solid but not critical safety improvement in the test wrapper.	Medium
	Always re-enable GC Prevent leaving GC disabled if exceptions occur by using a finally clause. This mirrors production safety to avoid memory/performance regressions during tests. tests/test_instrument_all_and_run.py [31-39] try: counter = time.perf_counter_ns() return_value = codeflash_wrapped(*args, **kwargs) codeflash_duration = time.perf_counter_ns() - counter except Exception as e: codeflash_duration = time.perf_counter_ns() - counter exception = e -gc.enable() +finally: + gc.enable() Suggestion importance[1-10]: 7 __ Why: Using a finally clause guarantees gc.enable() runs even if exceptions occur, enhancing test robustness with minimal risk and clear alignment with intent.	Medium
General	Sanitize ID components Guard against None values to avoid "None" strings in test_id and downstream tags. Use empty strings for optional fields like codeflash_test_class_name and ensure all components are strings. codeflash/code_utils/instrument_existing_tests.py [366-377] value=ast.JoinedStr( values=[ - ast.FormattedValue(value=ast.Name(id="codeflash_test_module_name", ctx=ast.Load()), conversion=-1), + ast.FormattedValue(value=ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[ast.Name(id="codeflash_test_module_name", ctx=ast.Load())], keywords=[]), conversion=-1), ast.Constant(value=":"), - ast.FormattedValue(value=ast.Name(id="codeflash_test_class_name", ctx=ast.Load()), conversion=-1), + ast.FormattedValue(value=ast.IfExp(test=ast.Name(id="codeflash_test_class_name", ctx=ast.Load()), body=ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[ast.Name(id="codeflash_test_class_name", ctx=ast.Load())], keywords=[]), orelse=ast.Constant(value=""),), conversion=-1), ast.Constant(value=":"), - ast.FormattedValue(value=ast.Name(id="codeflash_test_name", ctx=ast.Load()), conversion=-1), + ast.FormattedValue(value=ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[ast.Name(id="codeflash_test_name", ctx=ast.Load())], keywords=[]), conversion=-1), ast.Constant(value=":"), - ast.FormattedValue(value=ast.Name(id="codeflash_line_id", ctx=ast.Load()), conversion=-1), + ast.FormattedValue(value=ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[ast.Name(id="codeflash_line_id", ctx=ast.Load())], keywords=[]), conversion=-1), ast.Constant(value=":"), - ast.FormattedValue(value=ast.Name(id="codeflash_loop_index", ctx=ast.Load()), conversion=-1), + ast.FormattedValue(value=ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[ast.Name(id="codeflash_loop_index", ctx=ast.Load())], keywords=[]), conversion=-1), ] ) Suggestion importance[1-10]: 5 __ Why: The change avoids "None" appearing in test_id and ensures string conversion, which is a reasonable robustness improvement; however, it's not strictly required by the PR and slightly deviates from existing patterns used elsewhere.	Low

[codeflash-ai]

This PR is now faster! 🚀 Kevin Turcios accepted my code suggestion above.