Fix more edge cases with multiple imports and nested list comprehensions

Author: masenf

reflex/vars/dep_tracking.py

@@ -396,7 +396,14 @@ def _populate_dependencies(self) -> None:
             elif self.scan_status == ScanStatus.GETTING_VAR:
                 self.handle_getting_var(instruction)
             elif (
-                instruction.opname in ("LOAD_FAST", "LOAD_DEREF", "LOAD_FAST_BORROW")
+                instruction.opname
+                in (
+                    "LOAD_FAST",
+                    "LOAD_DEREF",
+                    "LOAD_FAST_BORROW",
+                    "LOAD_FAST_CHECK",
+                    "LOAD_FAST_AND_CLEAR",
+                )
                 and instruction.argval in self.tracked_locals
             ):
                 # bytecode loaded the class instance to the top of stack, next load instruction
@@ -405,7 +412,11 @@ def _populate_dependencies(self) -> None:
                 self.scan_status = ScanStatus.GETTING_ATTR
             elif (
                 instruction.opname
-                in ("LOAD_FAST_LOAD_FAST", "LOAD_FAST_BORROW_LOAD_FAST_BORROW")
+                in (
+                    "LOAD_FAST_LOAD_FAST",
+                    "LOAD_FAST_BORROW_LOAD_FAST_BORROW",
+                    "STORE_FAST_LOAD_FAST",
+                )
                 and instruction.argval[-1] in self.tracked_locals
             ):
                 # Double LOAD_FAST family instructions load multiple values onto the stack,
@@ -457,7 +468,6 @@ def _populate_dependencies(self) -> None:
                 )
                 # If we see a STORE_FAST, we can assign the top of stack to an aliased name.
                 self.top_of_stack = instruction.argval
-                self._last_import_name = None
             elif instruction.opname == "STORE_FAST" and self.top_of_stack is not None:
                 self.tracked_locals[instruction.argval] = self.tracked_locals.pop(
                     self.top_of_stack

tests/units/vars/test_dep_tracking.py

@@ -23,6 +23,7 @@ class DependencyTestState(State):
     count: rx.Field[int] = rx.field(default=0)
     name: rx.Field[str] = rx.field(default="test")
     items: rx.Field[list[str]] = rx.field(default_factory=list)
+    board: rx.Field[list[list[int]]] = rx.field(default_factory=list)
 
 
 class AnotherTestState(State):
@@ -102,6 +103,18 @@ def func_with_comprehension(self: DependencyTestState):
     assert tracker.dependencies == expected_deps
 
 
+def test_list_comprehension_dependencies_2():
+    """Test tracking dependencies in list comprehensions."""
+
+    def func_with_comprehension(self: DependencyTestState):
+        return [[self.board[r][c] for r in range(3)] for c in range(5)]
+
+    tracker = DependencyTracker(func_with_comprehension, DependencyTestState)
+
+    expected_deps = {DependencyTestState.get_full_name(): {"board"}}
+    assert tracker.dependencies == expected_deps
+
+
 def test_invalid_attribute_access():
     """Test that accessing invalid attributes raises VarValueError."""
 
@@ -167,6 +180,24 @@ async def get_state_import_from(self: DependencyTestState):
     assert tracker.dependencies == expected_deps
 
 
+def test_get_state_with_import_from_multiple():
+    """Test that get_state with function-local `from ... import ...` finds correct dependency."""
+
+    async def get_state_import_from(self: DependencyTestState):
+        from tests.units.states.upload import ChildFileUploadState, SubUploadState
+
+        return (await self.get_state(SubUploadState)).img, (
+            await self.get_state(ChildFileUploadState)
+        ).img_list
+
+    tracker = DependencyTracker(get_state_import_from, DependencyTestState)
+    expected_deps = {
+        tus_upload.SubUploadState.get_full_name(): {"img"},
+        tus_upload.ChildFileUploadState.get_full_name(): {"img_list"},
+    }
+    assert tracker.dependencies == expected_deps
+
+
 def test_get_state_with_import_from_as():
     """Test that get_state with function-local `from ... import ... as ...` finds correct dependency."""