fix Language server does not expose read/write metadata for symbols #3246

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
+                        },
+                    ),
                 }),
         ))
     }

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,
             },
         }
     }

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(),
+    );
+}

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();

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;