Add documentation hover for assignments with docstring

pyrefly/lib/binding/binding.rs

@@ -1973,6 +1973,9 @@ pub struct NameAssign {
     /// The Definition idx for this NameAssign, if infer_with_first_use is enabled.
     /// Used at solve time for inline first-use pinning and partial answer storage.
     pub def_idx: Option<Idx<Key>>,
+    /// If a standalone string literal immediately follows this assignment (PEP 257-style
+    /// variable docstring), this holds its range.
+    pub docstring_range: Option<TextRange>,
 }
 
 /// Data for a type alias binding.

pyrefly/lib/binding/bindings.rs

@@ -257,6 +257,10 @@ pub struct BindingsBuilder<'a> {
     lambda_yield_keys: Vec<(TextRange, Box<[Idx<KeyYield>]>, Box<[Idx<KeyYieldFrom>]>)>,
     /// See `BindingsInner::subsequently_initialized`.
     subsequently_initialized: SmallSet<Idx<KeyAnnotation>>,
+    /// The range of a PEP 257-style variable docstring (a standalone string literal
+    /// immediately following the current assignment statement), if any. Set by `stmts`
+    /// with look-ahead before calling `stmt`, consumed by `bind_single_name_assign`.
+    pub pending_var_docstring: Option<TextRange>,
 }
 
 /// An enum tracking whether we are in a generator expression
@@ -512,6 +516,7 @@ impl Bindings {
             deferred_bound_names: Vec::new(),
             lambda_yield_keys: Vec::new(),
             subsequently_initialized: SmallSet::new(),
+            pending_var_docstring: None,
         };
         builder.init_static_scope(&x.body, true);
         if module_info.name() != ModuleName::builtins() {
@@ -940,8 +945,20 @@ impl<'a> BindingsBuilder<'a> {
     }
 
     pub fn stmts(&mut self, xs: Vec<Stmt>, parent: &NestingContext) {
-        for x in xs {
+        let mut iter = xs.into_iter().peekable();
+        while let Some(x) = iter.next() {
+            // PEP 257-style variable docstrings: a standalone string literal immediately
+            // following an assignment is the variable's docstring. Set it before `stmt`
+            // so `bind_single_name_assign` can read it; clear it after.
+            if let Stmt::Assign(_) | Stmt::AnnAssign(_) = x
+                && let Some(Stmt::Expr(e)) = iter.peek()
+                && let Expr::StringLiteral(_) = e.value.as_ref()
+            {
+                self.pending_var_docstring = Some(e.range());
+            }
+
             self.stmt(x, parent);
+            self.pending_var_docstring = None;
         }
     }
 

pyrefly/lib/binding/target.rs

@@ -544,6 +544,7 @@ impl<'a> BindingsBuilder<'a> {
                 is_in_function_scope: self.scopes.in_function_scope(),
                 first_use: FirstUse::Undetermined,
                 def_idx: if uses_first_use { Some(def_idx) } else { None },
+                docstring_range: self.pending_var_docstring,
             }))
         };
         self.insert_binding_idx(def_idx, binding);

pyrefly/lib/export/definitions.rs

@@ -405,8 +405,37 @@ impl Definitions {
 
 impl<'a> DefinitionsBuilder<'a> {
     fn stmts(&mut self, xs: &[Stmt]) {
-        for x in xs {
+        for (i, x) in xs.iter().enumerate() {
             self.stmt(x);
+            // PEP 257-style variable docstrings: a standalone string literal expression
+            // immediately following an assignment statement is treated as the variable's docstring.
+            if let Stmt::Assign(_) | Stmt::AnnAssign(_) = x
+                && let Some(Stmt::Expr(e)) = xs.get(i + 1)
+                && let Expr::StringLiteral(_) = e.value.as_ref()
+            {
+                self.set_var_docstring(x, e.range());
+            }
+        }
+    }
+
+    /// Set `docstring_range` on all names defined by `assign_stmt`.
+    fn set_var_docstring(&mut self, assign_stmt: &Stmt, docstring_range: TextRange) {
+        let mut names = Vec::new();
+        match assign_stmt {
+            Stmt::Assign(x) => {
+                for t in &x.targets {
+                    Ast::expr_lvalue(t, &mut |n: &ExprName| names.push(n.id.clone()));
+                }
+            }
+            Stmt::AnnAssign(x) => {
+                Ast::expr_lvalue(&x.target, &mut |n: &ExprName| names.push(n.id.clone()));
+            }
+            _ => {}
+        }
+        for name in names {
+            if let Some(def) = self.inner.definitions.get_mut(&name) {
+                def.docstring_range = Some(docstring_range);
+            }
         }
     }
 

pyrefly/lib/state/ide.rs

@@ -209,6 +209,16 @@ fn create_intermediate_definition_from(
                     })),
                 };
             }
+            Binding::NameAssign(x) => {
+                return Some(IntermediateDefinition::Local(Export {
+                    location: def_key.range(),
+                    symbol_kind: current_binding.symbol_kind(),
+                    docstring_range: x.docstring_range,
+                    deprecation: None,
+                    is_final: false,
+                    special_export: None,
+                }));
+            }
             _ => {
                 return Some(IntermediateDefinition::Local(Export {
                     location: def_key.range(),

pyrefly/lib/test/lsp/hover_docstring.rs

@@ -351,6 +351,68 @@ Docstring Result: `Test docstring`
     );
 }
 
+#[test]
+fn module_variable_test() {
+    let code = r#"
+VAR = "abc"
+"""Some documentation."""
+print(VAR)
+#     ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
+    assert_eq!(
+        r#"
+# main.py
+4 | print(VAR)
+          ^
+Docstring Result: `Some documentation.`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn module_variable_itself_test() {
+    let code = r#"
+VAR = "abc"
+#^
+"""Some documentation."""
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
+    assert_eq!(
+        r#"
+# main.py
+2 | VAR = "abc"
+     ^
+Docstring Result: `Some documentation.`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn annotated_variable_test() {
+    let code = r#"
+VAR: str = "abc"
+"""Some documentation."""
+print(VAR)
+#     ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
+    assert_eq!(
+        r#"
+# main.py
+4 | print(VAR)
+          ^
+Docstring Result: `Some documentation.`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
 #[test]
 fn module_binding_test() {
     let lib = r#"
@@ -373,6 +435,78 @@ print(lib)
 Docstring Result: `Test docstring`
 
 
+# lib.py
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn variable_docstring_test() {
+    let code = r#"
+VAR = "abc"
+"""Some documentation."""
+print(VAR)
+#     ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
+    assert_eq!(
+        r#"
+# main.py
+4 | print(VAR)
+          ^
+Docstring Result: `Some documentation.`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn annotated_variable_docstring_test() {
+    let code = r#"
+VAR: str = "abc"
+"""Some documentation."""
+print(VAR)
+#     ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], test_report_factory(code));
+    assert_eq!(
+        r#"
+# main.py
+4 | print(VAR)
+          ^
+Docstring Result: `Some documentation.`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn cross_module_variable_docstring_test() {
+    let lib = r#"
+VAR = "abc"
+"""Some documentation."""
+"#;
+    let code = r#"
+from lib import VAR
+print(VAR)
+#     ^
+"#;
+    let report = get_batched_lsp_operations_report(
+        &[("main", code), ("lib", lib)],
+        test_report_factory(lib),
+    );
+    assert_eq!(
+        r#"
+# main.py
+3 | print(VAR)
+          ^
+Docstring Result: `Some documentation.`
+
+
 # lib.py
 "#
         .trim(),