fix inconsistent syntax highlighting for field of type inspect._ParameterKind #4751

pyrefly/lib/state/semantic_tokens.rs

@@ -263,24 +263,28 @@ fn range_overlaps(limit_range: Option<TextRange>, range: TextRange) -> bool {
 /// Classify an attribute's resolved type into a semantic token kind. For a union,
 /// every member must agree on the same kind; any disagreement (or a member that is
 /// a plain attribute) falls back to `PROPERTY`.
-fn attribute_semantic_token_type(ty: Type) -> SemanticTokenType {
+fn attribute_semantic_token_type(ty: Type, accessed_on_class: bool) -> SemanticTokenType {
     match ty {
         Type::Union(union) => {
             let mut members = union.members.into_iter();
             let Some(first) = members.next() else {
                 return SemanticTokenType::PROPERTY;
             };
-            let kind = attribute_semantic_token_type(first);
+            let kind = attribute_semantic_token_type(first, accessed_on_class);
             if kind == SemanticTokenType::PROPERTY {
                 return SemanticTokenType::PROPERTY;
             }
-            if members.all(|member| attribute_semantic_token_type(member) == kind) {
+            if members
+                .all(|member| attribute_semantic_token_type(member, accessed_on_class) == kind)
+            {
                 kind
             } else {
                 SemanticTokenType::PROPERTY
             }
         }
-        Type::Literal(lit) if matches!(lit.value, Lit::Enum(_)) => SemanticTokenType::ENUM_MEMBER,
+        Type::Literal(lit) if accessed_on_class && matches!(lit.value, Lit::Enum(_)) => {
+            SemanticTokenType::ENUM_MEMBER
+        }
         _ => {
             attribute_symbol_kind_from_type(&ty)
                 .to_lsp_semantic_token_type_with_modifiers()
@@ -421,8 +425,10 @@ impl SemanticTokenBuilder {
         get_type_of_attribute: &dyn Fn(TextRange) -> Option<Type>,
         get_symbol_kind: &dyn Fn(&Key) -> Option<(ModuleName, SymbolKind)>,
     ) {
+        let accessed_on_class = get_type_of_attribute(attr.value.range())
+            .is_some_and(|ty| matches!(ty, Type::ClassDef(_) | Type::Type(_)));
         let kind = get_type_of_attribute(attr.range())
-            .map(attribute_semantic_token_type)
+            .map(|ty| attribute_semantic_token_type(ty, accessed_on_class))
             .unwrap_or(SemanticTokenType::PROPERTY);
         self.push_if_in_range(attr.attr.range(), kind, Vec::new());
         attr.value

pyrefly/lib/test/lsp/semantic_tokens.rs

@@ -928,6 +928,91 @@ token-type: enumMember
     );
 }
 
+#[test]
+fn narrowed_enum_instance_attribute_test() {
+    let code = r#"
+from enum import Enum
+
+class E(Enum):
+    A = 1
+    B = 2
+
+class Holder:
+    kind: E
+
+def f(holder: Holder) -> None:
+    if holder.kind is E.A:
+        pass
+    elif holder.kind is E.B:
+        pass
+"#;
+    assert_full_semantic_tokens(
+        &[("main", code)],
+        r#"
+# main.py
+line: 1, column: 5, length: 4, text: enum
+token-type: namespace
+
+line: 1, column: 17, length: 4, text: Enum
+token-type: class
+
+line: 3, column: 6, length: 1, text: E
+token-type: class
+
+line: 3, column: 8, length: 4, text: Enum
+token-type: class
+
+line: 4, column: 4, length: 1, text: A
+token-type: variable, token-modifiers: [readonly]
+
+line: 5, column: 4, length: 1, text: B
+token-type: variable, token-modifiers: [readonly]
+
+line: 7, column: 6, length: 6, text: Holder
+token-type: class
+
+line: 8, column: 4, length: 4, text: kind
+token-type: variable
+
+line: 8, column: 10, length: 1, text: E
+token-type: class
+
+line: 10, column: 4, length: 1, text: f
+token-type: function
+
+line: 10, column: 6, length: 6, text: holder
+token-type: parameter
+
+line: 10, column: 14, length: 6, text: Holder
+token-type: class
+
+line: 11, column: 7, length: 6, text: holder
+token-type: parameter
+
+line: 11, column: 14, length: 4, text: kind
+token-type: property
+
+line: 11, column: 22, length: 1, text: E
+token-type: class
+
+line: 11, column: 24, length: 1, text: A
+token-type: enumMember
+
+line: 13, column: 9, length: 6, text: holder
+token-type: parameter
+
+line: 13, column: 16, length: 4, text: kind
+token-type: property
+
+line: 13, column: 24, length: 1, text: E
+token-type: class
+
+line: 13, column: 26, length: 1, text: B
+token-type: enumMember
+"#,
+    );
+}
+
 #[test]
 fn type_alias_test() {
     let code = r#"