diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index d1be83be33..3c6f058a09 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -67,6 +67,7 @@ use lsp_types::DidChangeWorkspaceFoldersParams; use lsp_types::DocumentDiagnosticParams; use lsp_types::DocumentDiagnosticReport; use lsp_types::DocumentHighlight; +use lsp_types::DocumentHighlightKind; use lsp_types::DocumentHighlightParams; use lsp_types::DocumentSymbol; use lsp_types::DocumentSymbolParams; @@ -4650,7 +4651,18 @@ impl Server { .find_local_references(&handle, position, true) .into_map(|range| DocumentHighlight { range: info.to_lsp_range(range), - kind: None, + kind: Some( + if transaction + .identifier_at(&handle, range.start()) + .expect("local references should point at identifiers") + .context + .is_write() + { + DocumentHighlightKind::WRITE + } else { + DocumentHighlightKind::READ + }, + ), }), )) } diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 65e4178e35..783d48acf0 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -286,6 +286,8 @@ pub(crate) enum IdentifierContext { base_range: TextRange, /// The range of the entire expression. range: TextRange, + /// Whether the attribute is being loaded, assigned to, or deleted. + expr_context: ExprContext, }, /// An identifier appeared as the name of a keyword argument. /// ex: `x` in `f(x=1)`. We also store some info about the callee `f` so @@ -341,6 +343,28 @@ pub(crate) enum IdentifierContext { MutableCapture, } +impl IdentifierContext { + pub(crate) fn is_write(&self) -> bool { + matches!( + self, + IdentifierContext::Expr(ExprContext::Store | ExprContext::Del) + | IdentifierContext::Attribute { + expr_context: ExprContext::Store | ExprContext::Del, + .. + } + | IdentifierContext::ImportedModule { .. } + | IdentifierContext::ImportedName { .. } + | IdentifierContext::FunctionDef { .. } + | IdentifierContext::MethodDef { .. } + | IdentifierContext::ClassDef { .. } + | IdentifierContext::Parameter + | IdentifierContext::TypeParameter + | IdentifierContext::ExceptionHandler + | IdentifierContext::PatternMatch(_) + ) + } +} + #[derive(Debug)] pub(crate) struct IdentifierWithContext { pub(crate) identifier: Identifier, @@ -495,6 +519,7 @@ impl IdentifierWithContext { context: IdentifierContext::Attribute { base_range: attr.value.range(), range: attr.range(), + expr_context: attr.ctx, }, } } diff --git a/pyrefly/lib/test/lsp/document_highlight.rs b/pyrefly/lib/test/lsp/document_highlight.rs new file mode 100644 index 0000000000..37d9069349 --- /dev/null +++ b/pyrefly/lib/test/lsp/document_highlight.rs @@ -0,0 +1,73 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use itertools::Itertools; +use lsp_types::DocumentHighlightKind; +use pretty_assertions::assert_eq; +use pyrefly_build::handle::Handle; +use ruff_text_size::TextSize; + +use crate::state::state::State; +use crate::test::util::code_frame_of_source_at_range; +use crate::test::util::get_batched_lsp_operations_report; + +fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String { + let transaction = state.transaction(); + let module_info = transaction.get_module_info(handle).unwrap(); + let highlights = transaction + .find_local_references(handle, position, true) + .into_iter() + .map(|range| { + let kind = if transaction + .identifier_at(handle, range.start()) + .expect("local references should point at identifiers") + .context + .is_write() + { + DocumentHighlightKind::WRITE + } else { + DocumentHighlightKind::READ + }; + format!( + "{}:\n{}", + if kind == DocumentHighlightKind::WRITE { + "DocumentHighlightKind::WRITE" + } else { + "DocumentHighlightKind::READ" + }, + code_frame_of_source_at_range(module_info.contents(), range) + ) + }) + .join("\n"); + format!("Highlights:\n{highlights}") +} + +#[test] +fn document_highlight_includes_read_write_kind() { + let code = r#" +x = 1 +y = x +# ^ +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +3 | y = x + ^ +Highlights: +DocumentHighlightKind::WRITE: +2 | x = 1 + ^ +DocumentHighlightKind::READ: +3 | y = x + ^ +"# + .trim(), + report.trim(), + ); +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/notebook_document_highlight.rs b/pyrefly/lib/test/lsp/lsp_interaction/notebook_document_highlight.rs index 38fb69a63c..e66e849683 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/notebook_document_highlight.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/notebook_document_highlight.rs @@ -5,6 +5,7 @@ * LICENSE file in the root directory of this source tree. */ +use lsp_types::DocumentHighlightKind; use serde_json::json; use crate::object_model::InitializeSettings; @@ -32,13 +33,15 @@ fn test_notebook_document_highlight() { "range": { "start": { "line": 0, "character": 0 }, "end": { "line": 0, "character": 1 } - } + }, + "kind": DocumentHighlightKind::WRITE }, { "range": { "start": { "line": 1, "character": 4 }, "end": { "line": 1, "character": 5 } - } + }, + "kind": DocumentHighlightKind::READ } ])) .unwrap(); diff --git a/pyrefly/lib/test/lsp/mod.rs b/pyrefly/lib/test/lsp/mod.rs index 08eb218850..8054fcbe9e 100644 --- a/pyrefly/lib/test/lsp/mod.rs +++ b/pyrefly/lib/test/lsp/mod.rs @@ -13,6 +13,7 @@ mod completion; mod declaration; mod definition; mod diagnostic; +mod document_highlight; mod document_symbols; mod folding_ranges; mod hover;