fix Inlay hint for attribute does not match inherited annotation #846

pyrefly/lib/lsp/wasm/inlay_hints.rs

@@ -342,10 +342,13 @@ impl<'a> Transaction<'a> {
                         _ => {}
                     }
                 }
-                key @ Key::Definition(_)
+                key @ Key::Definition(identifier)
                     if inlay_hint_config.variable_types
-                        && let Some(mut ty) = answers.get_type_at(idx) =>
+                        && let Some(ty) = answers.get_type_at(idx) =>
                 {
+                    let mut ty = self
+                        .class_field_type_at_definition(handle, *identifier)
+                        .unwrap_or(ty);
                     let mut insertable = true;
                     if let Some(constructor) = self.new_type_constructor_signature(handle, &ty) {
                         ty = constructor;

pyrefly/lib/state/lsp.rs

@@ -81,6 +81,7 @@ use crate::alt::attr::AttrInfo;
 use crate::binding::binding::Binding;
 use crate::binding::binding::Key;
 use crate::binding::binding::KeyAnnotation;
+use crate::binding::binding::KeyClassField;
 use crate::config::error_kind::ErrorKind;
 use crate::error::suppress::detect_line_ending;
 use crate::export::exports::Export;
@@ -895,6 +896,43 @@ impl<'a> Transaction<'a> {
         answers.get_annotation_type_at(idx)
     }
 
+    /// Returns the resolved declaration type without applying attribute-access substitutions.
+    pub(crate) fn class_field_type_at_definition(
+        &self,
+        handle: &Handle,
+        identifier: ShortIdentifier,
+    ) -> Option<Type> {
+        let ast = self.get_ast(handle)?;
+        let mut enclosing_class = None;
+        for node in Ast::locate_node(&ast, identifier.range().start())
+            .into_iter()
+            .skip(1)
+        {
+            match node {
+                AnyNodeRef::StmtFunctionDef(_)
+                | AnyNodeRef::ExprLambda(_)
+                | AnyNodeRef::ExprListComp(_)
+                | AnyNodeRef::ExprSetComp(_)
+                | AnyNodeRef::ExprDictComp(_)
+                | AnyNodeRef::ExprGenerator(_) => return None,
+                AnyNodeRef::StmtClassDef(class_def) => {
+                    enclosing_class = Some(class_def);
+                    break;
+                }
+                _ => {}
+            }
+        }
+        let enclosing_class = enclosing_class?;
+        let answers = self.get_answers(handle)?;
+        let bindings = answers.bindings();
+        let class_index = bindings.class_def_index(enclosing_class)?;
+        let name = Name::new(bindings.module().code_at(identifier.range()));
+        let key = KeyClassField(class_index, name);
+        let idx = bindings.key_to_idx_hashed_opt(Hashed::new(&key))?;
+        let field = answers.get_idx(idx)?;
+        field.has_explicit_annotation().then(|| field.ty())
+    }
+
     fn get_chosen_overload_trace_for_surface(
         &self,
         handle: &Handle,
@@ -1507,7 +1545,12 @@ impl<'a> Transaction<'a> {
                 if !bindings.is_valid_key(&key) {
                     return None;
                 }
-                let mut ty = answers.get_type_at(bindings.key_to_idx(&key))?;
+                let mut ty = if matches!(context, IdentifierContext::Expr(ExprContext::Store)) {
+                    self.class_field_type_at_definition(handle, ShortIdentifier::new(identifier))
+                        .or_else(|| answers.get_type_at(bindings.key_to_idx(&key)))?
+                } else {
+                    answers.get_type_at(bindings.key_to_idx(&key))?
+                };
                 // Only a plain expression reference coerces to its callee signature.
                 if coerce_callees && let IdentifierContext::Expr(_) = context {
                     let call_args_range = self.callee_at(handle, position).and_then(

pyrefly/lib/test/lsp/hover_type.rs

@@ -280,6 +280,77 @@ Hover Result: `Literal[2]`
     );
 }
 
+#[test]
+fn class_field_comprehension_shadowing_hover() {
+    let code = r#"
+class C:
+    x: str = "class field"
+    a = [x for x in (1, 2)]
+#              ^
+    b = {x for x in (1, 2)}
+#              ^
+    c = {x: x for x in (1, 2)}
+#                 ^
+    d = (x for x in (1, 2))
+#              ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+4 |     a = [x for x in (1, 2)]
+                   ^
+Hover Result: `Literal[1, 2]`
+
+6 |     b = {x for x in (1, 2)}
+                   ^
+Hover Result: `Literal[1, 2]`
+
+8 |     c = {x: x for x in (1, 2)}
+                      ^
+Hover Result: `Literal[1, 2]`
+
+10 |     d = (x for x in (1, 2))
+                    ^
+Hover Result: `Literal[1, 2]`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn inherited_classvar_hover_prefers_resolved_field_type() {
+    let code = r#"
+from typing import ClassVar
+
+class A:
+    x: ClassVar[tuple[str, ...]] = ("A",)
+
+class B(A):
+    x = ("B",)
+#   ^
+
+reveal = B.x
+#          ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+8 |     x = ("B",)
+        ^
+Hover Result: `tuple[str, ...]`
+
+11 | reveal = B.x
+                ^
+Hover Result: `tuple[str, ...]`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
 #[test]
 fn import_test() {
     let code = r#"

pyrefly/lib/test/lsp/inlay_hint.rs

@@ -69,6 +69,31 @@ def f(xs: list[int]) -> None:
     );
 }
 
+#[test]
+fn inherited_classvar_inlay_hint_prefers_resolved_field_type() {
+    let code = r#"
+from typing import ClassVar
+
+def value() -> str:
+    return "value"
+
+class A:
+    x: ClassVar[tuple[str, ...]] = ("A",)
+
+class B(A):
+    x = (value(),)
+"#;
+    assert_eq!(
+        r#"
+# main.py
+11 |     x = (value(),)
+          ^ inlay-hint: `: tuple[str, ...]`
+"#
+        .trim(),
+        generate_inlay_hint_report(code, Default::default()).trim(),
+    );
+}
+
 #[test]
 fn basic_test() {
     let code = r#"from typing import Literal

pyrefly/lib/test/lsp/write_target_type.rs

@@ -23,6 +23,7 @@ use crate::state::require::Require;
 use crate::state::state::State;
 use crate::test::util::get_batched_lsp_operations_report_allow_error;
 use crate::test::util::mk_multi_file_state;
+use crate::test::util::mk_multi_file_state_assert_no_errors;
 
 fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String {
     let transaction = state.transaction();
@@ -38,6 +39,35 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String
     format!("Hover: {hover}\nComputed: {computed}")
 }
 
+#[test]
+fn class_field_declaration_preserves_self() {
+    let code = "from typing import Self\nclass C:\n    x: Self | None = None\n";
+    let (handles, state) =
+        mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Exports);
+    let handle = handles.get("main").unwrap();
+    let position = TextSize::try_from(code.find("x:").unwrap()).unwrap();
+    let transaction = state.transaction();
+    assert_eq!(
+        "Self@C | None",
+        transaction
+            .get_type_at(handle, position)
+            .unwrap()
+            .to_string(),
+    );
+    for range in [
+        TextRange::empty(position),
+        TextRange::at(position, 1.into()),
+    ] {
+        assert_eq!(
+            "Self@C | None",
+            transaction
+                .get_computed_type_at_range(handle, range)
+                .unwrap()
+                .to_string(),
+        );
+    }
+}
+
 /// A store target reports the type that is being assigned to it. When the value
 /// does not fit the attribute's declared type, the assignment recovers with the
 /// declared type and the target reports that.