Refactor: create_wrapper_function in instrument_existing_tests to reduce complexity and isolate wrapper-building concerns

[shreejaykurhade]

Background

The create_wrapper_function implementation in codeflash/code_utils/instrument_existing_tests.py had become a large monolithic AST builder responsible for several distinct concerns at the same time:

• invocation tracking
• stdout tag generation
• timing and exception handling
• framework device synchronization
• behavior-mode persistence
• wrapper argument construction

Because all of this logic lived in one function, the code was difficult to review, difficult to modify safely, and unnecessarily risky for future instrumentation work.

Previous Issue #526

The main issue was not functional breakage but maintainability risk.

The function had grown large enough that:

• small edits required understanding a long AST statement list end to end
• debugging generated wrapper output was harder than necessary
• future changes had a higher chance of introducing subtle ordering or behavior regressions
• code review signal was low because business intent was buried inside AST construction detail

This is the kind of function that tends to slow down future development even when it still works correctly today.

Improvement Made

The wrapper generation logic was refactored into smaller private helper builders, each with a single responsibility.

The refactor separates:

• perf counter call creation
• duration assignment creation
• invocation tracking statement generation
• stdout tag statement generation
• try/except timing block generation
• behavior-mode persistence generation
• wrapper argument generation

After the refactor, create_wrapper_function acts as a thin orchestrator that assembles the wrapper in the correct order.

Why This Is Better

This change improves the codebase in several ways:

• lowers cognitive load for future readers and reviewers
• makes wrapper generation easier to debug by concern
• reduces the blast radius of future changes
• preserves behavior while improving readability
• makes the instrumentation pipeline easier to extend safely

Scope

This was intentionally kept as a pure refactor.

There are no intended changes to:

• generated instrumentation behavior
• timing behavior
• exception flow
• behavior-mode sqlite persistence
• framework sync logic

Validation

The refactor was validated with focused instrumentation coverage:

• uv run pytest -q tests/test_inject_profiling_used_frameworks.py
• uv run ruff check codeflash/code_utils/instrument_existing_tests.py

These checks confirmed that the generated wrapper output remains correct for the covered framework and mode combinations.

Notes On Full Test Run

Some broader uv run pytest collection errors may still appear in this workspace, but they are unrelated to this refactor. They come from missing optional test dependencies in code_to_optimize test modules and are outside the scope of this change.

ERROR code_to_optimize/code_directories/retriever/import_test.py
ERROR code_to_optimize/tests/pytest/test_book_catalog.py
ERROR code_to_optimize/tests/pytest/test_book_catalog_2.py
ERROR code_to_optimize/tests/pytest/test_book_catalog_3.py
ERROR code_to_optimize/tests/pytest/test_jax_jit_code.py
ERROR code_to_optimize/tests/pytest/test_numba_jit_code.py
ERROR code_to_optimize/tests/pytest/test_torch_jit_code.py
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Interrupted: 9 errors during collection !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
=================================== 1 skipped, 2 warnings, 9 errors in 3.11s

Risk Assessment

Risk is low.

The change is structural, localized, and supported by output-oriented tests around instrumentation generation. Review attention should focus on statement ordering and equivalence of emitted wrapper code.

Files Changed

• codeflash/code_utils/instrument_existing_tests.py

[CLAassistant]

All committers have signed the CLA.

[shreejaykurhade]

@Saga4 @KRRT7 implemented Issue #524 and #525