diff --git a/pyrefly/lib/state/semantic_tokens.rs b/pyrefly/lib/state/semantic_tokens.rs index d50cf8484c..fc83dabc7f 100644 --- a/pyrefly/lib/state/semantic_tokens.rs +++ b/pyrefly/lib/state/semantic_tokens.rs @@ -263,24 +263,28 @@ fn range_overlaps(limit_range: Option, 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, 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 diff --git a/pyrefly/lib/test/lsp/semantic_tokens.rs b/pyrefly/lib/test/lsp/semantic_tokens.rs index 0afc5fd2a9..784ee5d51c 100644 --- a/pyrefly/lib/test/lsp/semantic_tokens.rs +++ b/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#"