Skip to content

Fix "Go to Definition" for variables used in list comprehensions(LSP highlighting not working for variables defined in list comprehension #2549)#2597

Open
pswitchy wants to merge 2 commits intofacebook:mainfrom
pswitchy:feat/lsp-highlighting
Open

Fix "Go to Definition" for variables used in list comprehensions(LSP highlighting not working for variables defined in list comprehension #2549)#2597
pswitchy wants to merge 2 commits intofacebook:mainfrom
pswitchy:feat/lsp-highlighting

Conversation

@pswitchy
Copy link

@pswitchy pswitchy commented Mar 1, 2026

Motivation:
Currently, when you are working with a list comprehension like [x for x in items], the "Go to Definition" feature works for the variable after the for (the target x), but it fails and returns nothing if you try it on the variable before the for (the usage x).

This PR fixes that bug so that clicking on the usage x will correctly take you to where it is defined.

Approach:
The bug happens because of how the Pyrefly language server tries to follow the chain of variable references.

  • For variables inside comprehensions, the chain goes through a special step called IterableValue.
  • The code that searches for definitions (find_definition_key_from in pyrefly/lib/state/ide.rs) did not know how to handle this IterableValue step, so it would give up and return nothing.
  • I updated the code to recognize IterableValue as a valid stopping point. Now, when the search hits this point, it correctly returns the location of the variable.

Tests:

  • I added two tests in pyrefly/lib/test/lsp/definition.rs to make sure "Go to Definition" works for both parts of a list comprehension and to prevent this bug from coming back.
  • I manually analyzed code bases and tested all the changes to make sure everything works correctly.

Tests Screenshots
Screenshot 2026-03-01 184413
Screenshot 2026-03-01 184357

@meta-codesync
Copy link

meta-codesync bot commented Mar 5, 2026

@yangdanny97 has imported this pull request. If you are a Meta employee, you can view this in D95438997.

@yangdanny97 yangdanny97 self-assigned this Mar 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants