This optimization achieves a **14% runtime improvement** (4.80ms → 4.19ms) through several targeted micro-optimizations that reduce overhead in hot code paths:
## Key Performance Improvements
### 1. **Eliminated Redundant Dictionary Lookups via Caching**
In `CodeStringsMarkdown` properties (`flat`, `file_to_path`), the original code called `self._cache.get("key")` twice per invocation. The optimized version caches the result in a local variable:
```python
# Before: two lookups
if self._cache.get("flat") is not None:
return self._cache["flat"]
# After: one lookup
cached = self._cache.get("flat")
if cached is not None:
return cached
```
This eliminates redundant hash table lookups in frequently accessed properties.
### 2. **Replaced `dict.setdefault()` for Atomic List Operations**
In `_analyze_imports_in_optimized_code`, the original code used an if-check followed by assignment for the helpers dictionary:
```python
# Before: check + assign (two operations)
if func_name in file_entry:
file_entry[func_name].append(helper)
else:
file_entry[func_name] = [helper]
# After: single atomic operation
helpers_by_file_and_func[module_name].setdefault(func_name, []).append(helper)
```
The `setdefault()` approach reduces the operation to a single dictionary call, eliminating the membership test.
### 3. **Hoisted `as_posix()` Calls Outside String Formatting**
In the `markdown` property, path conversion was moved outside the f-string:
```python
# Before: as_posix() called inside f-string
f"```python:{code_string.file_path.as_posix()}\n..."
# After: precomputed in conditional branch
if code_string.file_path:
file_path_str = code_string.file_path.as_posix()
result.append(f"```python:{file_path_str}\n...")
```
This avoids repeated method calls during string formatting.
### 4. **Optimized Set Membership Tests with Early Exit**
The most impactful change replaced `set.intersection()` with short-circuit boolean checks:
```python
# Before: creates intermediate set via intersection
is_called = bool(possible_call_names.intersection(called_function_names))
# After: early-exit on first match
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
```
With ~200 helpers in large-scale tests, this avoids creating temporary sets for every comparison, showing **50% speedup** in the large helper test (1.31ms → 868μs).
### 5. **Minimized Repeated Attribute Access**
Variables like `entrypoint_file_path`, `attr_name`, and `value_id` are now cached before use, reducing attribute lookups in the AST traversal loop.
## Impact Based on Test Results
- **Small workloads** (10-50 helpers): 10-16% speedup from reduced dict lookups
- **Large workloads** (200 helpers): 50% speedup due to eliminated set operations in the helper-checking loop
- **Edge cases** (syntax errors, missing functions): Minimal overhead, consistent 2-3% improvement
This optimization is particularly valuable when `detect_unused_helper_functions` is called repeatedly during code analysis pipelines, as the cumulative effect of these micro-optimizations scales with the number of helper functions and code blocks analyzed.
Consolidated prune_cst_for_read_only_code and prune_cst_for_testgen_code into
prune_cst_for_context with include_target_in_output and include_init_dunder flags.
Consolidated prune_cst_for_read_only_code and prune_cst_for_testgen_code into
prune_cst_for_context with include_target_in_output and include_init_dunder flags.
Assignments that don't reference module-level definitions are now placed
right after imports. Only assignments that reference classes/functions
are placed after those definitions to prevent NameError.
Update test_no_targets_found to expect outer class to be kept when
targeting nested class methods, and add test for nonexistent targets.
Update test_multi_file_replcement01 to expect global assignments at
module end rather than after imports.
When LLM-generated optimizations include module-level function calls like
`_register(MessageKind.ASK, ...)`, they were being inserted right after
imports, BEFORE the function definition they reference, causing NameError
at module load time.
Changes:
- Add GlobalStatementTransformer to append global statements at module end
- Reorder transformations: functions → assignments → statements
- Remove unused ImportInserter class
- Update test expectations to reflect new placement behavior
When LLM-generated optimizations include module-level code like
`_REIFIERS = {MessageKind.XXX: ...}`, the global assignment was being
inserted right after imports, BEFORE the class definition it referenced,
causing NameError at module load time.
Changes:
- GlobalAssignmentTransformer now inserts assignments after all
class/function definitions instead of right after imports
- GlobalStatementCollector now skips AnnAssign (annotated assignments)
so they are handled by GlobalAssignmentCollector instead
Classes used as dependencies (enums, dataclasses, types) were being
excluded from the optimization context even when marked as used by
the target function. This caused NameError when the LLM used these
types in generated optimizations.
Fix issue where enum/class attribute access like `MessageKind.VALUE` was not
tracking `MessageKind` as a dependency. The original code skipped all Names
inside Attribute nodes within classes, but this incorrectly filtered out
legitimate references.
Now properly distinguishes between:
- `.attr` part (e.g., `x` in `self.x`) - not tracked (attribute names)
- `.value` part (e.g., `MessageKind` in `MessageKind.VALUE`) - tracked
When instrumenting classes that inherit from Playwright's Page (like
SkyvernPage), the instance state contains async event loop references
including asyncgen_hooks which cannot be pickled. PicklePatcher gracefully
handles these by replacing unpicklable objects with placeholders.
Add get_external_base_class_inits to extract __init__ methods from external
library base classes (e.g., collections.UserDict) when project classes inherit
from them. This helps the LLM understand constructor signatures for mocking.
Add GlobalFunctionCollector and GlobalFunctionTransformer to collect and
insert module-level function definitions introduced by LLM optimizations.
This fixes NameError when optimized code introduces new helper functions
like @lru_cache decorated functions that are used by the optimized method.
GlobalAssignmentCollector only handled cst.Assign but not cst.AnnAssign
(annotated assignments like `X: int = 1`). When the LLM generated
optimizations with annotated module-level variables, these weren't
copied to the target file, causing NameError at runtime.
- Add visit_AnnAssign to GlobalAssignmentCollector
- Add leave_AnnAssign to GlobalAssignmentTransformer
- Update type hints to include cst.AnnAssign
- Add test for annotated assignment handling
Mock objects from unittest.mock can have corrupted internal state after
pickling, causing __repr__ to raise AttributeError. Added safe_repr
wrapper to gracefully handle these failures during test result comparison.
Support classes using __slots__ and C extension types (like Playwright's
Page) that don't have a __dict__ attribute. Previously raised ValueError,
now captures slot values or public non-callable attributes as fallback.
- Update test_get_imported_class_definitions_includes_dataclass_decorators
to expect both base class and derived class to be extracted
- Add test_get_imported_class_definitions_extracts_multilevel_inheritance
to verify multi-level inheritance chains are fully extracted
When extracting imported class definitions for testgen context, also
extract base classes from the same module. This ensures the full
inheritance chain is available for understanding constructor signatures.
For example, when LLMConfig inherits from LLMConfigBase, both classes
are now included in the context so the LLM can see all required
positional arguments from parent classes.
