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.
Show whether improvement is from runtime, throughput, or concurrency
in the display output (e.g., "45% runtime improvement" instead of
just "45% improvement").
