impl extract function, part of #746 #1786
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
asukaminato0721
force-pushed
the
746
branch
2 times, most recently
from
5e5ecd1 to
41161a7
Compare
There was a problem hiding this comment.
Pull request overview
This PR implements the "extract function" refactoring feature as part of issue #746. It adds the ability to select a block of Python code and extract it into a helper function, automatically inferring parameters and return values. The implementation includes a new LocalRefactorCodeAction type, visitor-based identifier collection, and integration with the LSP server to advertise the REFACTOR_EXTRACT capability.
Key Changes:
- Introduced
extract_function_code_actionsthat analyzes selected code blocks, identifies parameters (from loaded variables) and returns (from stored variables used later), and generates workspace edits to create a helper function and replace the selection with a function call - Added visitor-based identifier collection that distinguishes between loads, stores, and synthetic loads from augmented assignments
- Integrated extract function actions into the LSP server alongside existing quickfix actions
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
pyrefly/lib/test/lsp/lsp_interaction/basic.rs |
Updated server capabilities test to include "refactor.extract" in advertised code action kinds |
pyrefly/lib/test/lsp/code_actions.rs |
Added test infrastructure (helper functions) and basic integration test for extract function with augmented assignment |
pyrefly/lib/state/lsp.rs |
Core implementation including LocalRefactorCodeAction struct, extract_function_code_actions method, and helper functions for identifier collection, dedenting, indenting, and name generation |
pyrefly/lib/lsp/non_wasm/server.rs |
Integrated extract function actions into the code action handler and registered REFACTOR_EXTRACT capability |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
286
to
340
| #[test] | ||
| fn extract_function_basic_refactor() { | ||
| let code = r#" | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| let (handles, state) = | ||
| mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Everything); | ||
| let handle = handles.get("main").unwrap(); | ||
| let transaction = state.transaction(); | ||
| let module_info = transaction.get_module_info(handle).unwrap(); | ||
| let selection = find_marked_range(module_info.contents()); | ||
| let actions = transaction | ||
| .extract_function_code_actions(handle, selection) | ||
| .unwrap_or_default(); | ||
| assert!(!actions.is_empty(), "expected extract refactor action"); | ||
| let updated = apply_refactor_edits_for_module(&module_info, &actions[0].edits); | ||
| let expected = r#" | ||
| def extracted_function(item, total_sum): | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| return total_sum | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| total_sum = extracted_function(item, total_sum) | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| assert_eq!(expected.trim(), updated.trim()); | ||
| } |
There was a problem hiding this comment.
Test coverage for the extract function feature is limited. Only one basic scenario is tested. Consider adding tests for:
- Selections with no parameters or return values
- Multiple return values (tuple unpacking)
- Edge cases: empty selections, whitespace-only selections
- Rejection cases: selections containing return/break/continue/raise/function/class definitions
- Various indentation scenarios
- Synthetic loads from augmented assignments without regular loads
- Name collision scenarios for the generated function name
Contributor
There was a problem hiding this comment.
agreed. please add more tests, even if they are disabled
Contributor
There was a problem hiding this comment.
some other ones I can think of:
- function is within a class (but no classes in range), does the indent work?
Contributor
Author
|
this aim to be a basic impl, as copilot said, too many edge cases... |
kinto0
requested changes
Dec 9, 2025
Contributor
kinto0
left a comment
kinto0
left a comment
There was a problem hiding this comment.
awesome work! this is exciting. a few comments:
lsp.rsis much too big. can you put all the new additions in it's own file? maybe in alsp/quick_fixesfolder? (no need to move the other stuff, we can do that later)- a few comments about reuse and tests. I don't mind if the feature isn't implemented, I just don't want it crashing because of a crazy situation
Comment on lines
286
to
340
| #[test] | ||
| fn extract_function_basic_refactor() { | ||
| let code = r#" | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| let (handles, state) = | ||
| mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Everything); | ||
| let handle = handles.get("main").unwrap(); | ||
| let transaction = state.transaction(); | ||
| let module_info = transaction.get_module_info(handle).unwrap(); | ||
| let selection = find_marked_range(module_info.contents()); | ||
| let actions = transaction | ||
| .extract_function_code_actions(handle, selection) | ||
| .unwrap_or_default(); | ||
| assert!(!actions.is_empty(), "expected extract refactor action"); | ||
| let updated = apply_refactor_edits_for_module(&module_info, &actions[0].edits); | ||
| let expected = r#" | ||
| def extracted_function(item, total_sum): | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| return total_sum | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| total_sum = extracted_function(item, total_sum) | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| assert_eq!(expected.trim(), updated.trim()); | ||
| } |
Contributor
There was a problem hiding this comment.
agreed. please add more tests, even if they are disabled
pyrefly/lib/state/lsp.rs
Outdated
| found | ||
| } | ||
|
|
||
| fn find_enclosing_module_statement_range( |
Contributor
There was a problem hiding this comment.
@jvansch1 i think you needed something similar for call hierarchy. any way to reuse it?
Comment on lines
286
to
340
| #[test] | ||
| fn extract_function_basic_refactor() { | ||
| let code = r#" | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| let (handles, state) = | ||
| mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Everything); | ||
| let handle = handles.get("main").unwrap(); | ||
| let transaction = state.transaction(); | ||
| let module_info = transaction.get_module_info(handle).unwrap(); | ||
| let selection = find_marked_range(module_info.contents()); | ||
| let actions = transaction | ||
| .extract_function_code_actions(handle, selection) | ||
| .unwrap_or_default(); | ||
| assert!(!actions.is_empty(), "expected extract refactor action"); | ||
| let updated = apply_refactor_edits_for_module(&module_info, &actions[0].edits); | ||
| let expected = r#" | ||
| def extracted_function(item, total_sum): | ||
| squared_value = item * item | ||
| if squared_value > 100: | ||
| print(f"Large value detected: {squared_value}") | ||
| total_sum += squared_value | ||
| return total_sum | ||
| def process_data(data_list): | ||
| total_sum = 0 | ||
| for item in data_list: | ||
| # EXTRACT-START | ||
| total_sum = extracted_function(item, total_sum) | ||
| # EXTRACT-END | ||
| return total_sum | ||
| if __name__ == "__main__": | ||
| data = [1, 5, 12, 8, 15] | ||
| result = process_data(data) | ||
| print(f"The final sum is: {result}") | ||
| "#; | ||
| assert_eq!(expected.trim(), updated.trim()); | ||
| } |
Contributor
There was a problem hiding this comment.
some other ones I can think of:
- function is within a class (but no classes in range), does the indent work?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
part of #746
Introduced extract_function_code_actions plus LocalRefactorCodeAction plumbing so the state layer can produce workspace edits for a selected block; it now dedents the selection, synthesizes helper definitions/calls, infers parameters/returns (including augmented assignments), and rejects unsupported selections (returns/breaks/etc.).
Added a visitor-based identifier collector that walks statements/expressions to classify loads, stores, post-selection reads, and synthetic aug-assign loads, which drives the parameter/return heuristics.
Updated the non-wasm LSP server to advertise REFACTOR_EXTRACT, convert LocalRefactorCodeAction edits into URIs, and merge them alongside existing quick fixes.
Test Plan
Added an integration test that selects a block via markers, requests refactor actions, applies the returned edits, and asserts the helper/function call match the expected structure.