ENG-8509 Dependency tracking for locally imported states

CodSpeed Performance Report

Merging #6035 will not alter performance

_{Comparing masenf/var-dep-tracking-imports (24b1319) with main (bf49c88)}

Summary

✅ 8 untouched

Greptile Overview

Greptile Summary

This PR enhances the dependency tracking system to support function-local imports when using get_state in computed variables. The changes enable all common import patterns (from ... import ..., import ..., and aliased imports) to be properly tracked through bytecode analysis.

Key Changes

• New scan status: Added GETTING_STATE_POST_AWAIT to handle state access after awaiting get_state(), enabling both inline attribute access and variable assignment patterns
• Import tracking: Implemented IMPORT_NAME and IMPORT_FROM bytecode instruction handling to track modules and imported classes in the function's local scope
• Attribute chaining: Enhanced handle_getting_state to support LOAD_ATTR and LOAD_METHOD for accessing nested module attributes (e.g., module.SubModule.StateClass)
• Equality expressions: Added LOAD_FAST_LOAD_FAST family instruction support for tracking dependencies in comparison operations
• Better error handling: Introduced UntrackedLocalVarError for clearer errors when referencing untracked local variables
• Helper utilities: Added assert_base_state() and get_tracked_local() to validate and retrieve tracked objects

Test Coverage

The PR includes comprehensive test coverage for:

• Direct get_state() inline access without variable assignment
• All import variants: from X import Y, from X import Y as Z, import X, import X as Y
• Method calls on imported states
• Equality expressions with state attributes
• Error cases for untracked variables

The implementation correctly handles the complexity of Python's import bytecode while maintaining backward compatibility.

Confidence Score: 4/5

• This PR is safe to merge with low risk - it addresses a specific limitation in dependency tracking with comprehensive test coverage.
• The implementation is solid with thorough testing, but the bytecode analysis complexity and dynamic import execution introduce some risk. The code handles multiple edge cases correctly (aliased imports, nested modules, inline vs assigned state access). The score is 4 rather than 5 due to the inherent complexity of bytecode manipulation and potential for Python version-specific bytecode variations.
• No files require special attention - all changes are well-structured with appropriate test coverage.

Important Files Changed

File Analysis

Sequence Diagram

sequenceDiagram
    participant User as User Code (async function)
    participant DT as DependencyTracker
    participant Bytecode as Python Bytecode
    participant Imports as importlib/sys.modules
    participant State as BaseState

    User->>DT: __post_init__()
    DT->>DT: Initialize tracked_locals with self
    DT->>DT: _populate_dependencies()
    
    loop For each bytecode instruction
        DT->>Bytecode: Get instruction
        
        alt IMPORT_NAME
            Bytecode->>DT: instruction.argval = module name
            DT->>Imports: importlib.import_module(module)
            Imports-->>DT: Module loaded
            DT->>DT: Track module in tracked_locals
        
        else IMPORT_FROM
            Bytecode->>DT: instruction.argval = imported name
            DT->>Imports: Get attribute from module
            Imports-->>DT: Return imported object
            DT->>DT: Track imported object in tracked_locals
        
        else LOAD_FAST (tracked local)
            Bytecode->>DT: Load tracked local to stack
            DT->>DT: Set top_of_stack, status=GETTING_ATTR
        
        else LOAD_ATTR on get_state
            Bytecode->>DT: Call get_state
            DT->>DT: Set status=GETTING_STATE
            DT->>DT: Resolve state class from tracked_locals
        
        else END_SEND (await complete)
            Bytecode->>DT: Await finished
            DT->>DT: Set status=GETTING_STATE_POST_AWAIT
            DT->>DT: Track state instance
        
        else LOAD_ATTR on state
            Bytecode->>DT: Access state.attribute
            DT->>State: Get attribute
            State-->>DT: Attribute info
            DT->>DT: Record dependency
        
        else LOAD_FAST_LOAD_FAST
            Bytecode->>DT: Load multiple values (equality check)
            DT->>DT: Track rightmost value as top_of_stack
        end
    end
    
    DT-->>User: Dependencies collected