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

crates/pyrefly_types/src/types.rs

@@ -12,6 +12,7 @@ use std::fmt::Display;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::mem;
+use std::slice;
 use std::sync::Arc;
 
 use dupe::Dupe;
@@ -2436,6 +2437,14 @@ impl Type {
         }
     }
 
+    /// Borrow the top-level union members, or a singleton slice for a non-union type.
+    pub fn as_unions(&self) -> &[Type] {
+        match self {
+            Type::Union(union) => &union.members,
+            _ => slice::from_ref(self),
+        }
+    }
+
     pub fn into_unions(self) -> Vec<Type> {
         match self {
             Type::Union(u) => u.members,

pyrefly/benches/README.md

@@ -2,10 +2,10 @@
 
 Real-world benchmarks over a large, pinned PyTorch checkout (15k+ Python files)
 across all cores. They cover interactive LSP latency (`cold_start`,
-`error_propagation`, `workspace_symbol`), cold whole-project batch throughput
-(`full_check`), and cold whole-project indexing (`indexed_memory`). A separate
-`pytorch_memory` target reports the memory that index holds, for the reason in
-its section below.
+`error_propagation`, `workspace_symbol`), cold whole-project
+batch throughput (`full_check`), and cold whole-project indexing (`indexed_memory`).
+A separate `pytorch_memory` target reports the memory that index holds, for the
+reason in its section below.
 
 For the full command reference (all flags, micro benchmarks, cargo/buck forms),
 see `.agents/skills/benchmark-pyrefly/SKILL.md`.

pyrefly/benches/pytorch/main.rs

@@ -7,9 +7,9 @@
 
 //! PyTorch walltime benchmarks. Each benchmark lives in its own module
 //! (`cold_start`, `error_propagation`, `full_check`, `indexed_memory`,
-//! `workspace_symbol`) sharing the checkout harness in [`common`]; this crate root
-//! just aggregates their criterion groups into one binary, so a single
-//! `pytorch_bench` target builds and runs all of them.
+//! `workspace_symbol`) sharing the checkout harness in
+//! [`common`]. This crate root aggregates their Criterion groups into one binary.
+//! A single `pytorch_bench` target builds and runs all of them.
 //! Individual benchmarks are still selectable by name at runtime, e.g.
 //! `cargo bench -p pyrefly --bench pytorch -- cold_start`.
 

pyrefly/lib/lsp/wasm/semantic_tokens.rs

@@ -6,18 +6,64 @@
  */
 
 use lsp_types::SemanticToken;
+use lsp_types::SemanticTokenTypes;
 use pyrefly_build::handle::Handle;
+use pyrefly_types::literal::Lit;
+use pyrefly_types::types::Type;
+use ruff_python_ast::ExprAttribute;
+use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 
 use crate::binding::binding::Key;
 use crate::state::lsp::FindPreference;
 use crate::state::lsp::ImportBehavior;
+use crate::state::lsp::attribute_symbol_kind_from_type;
 use crate::state::semantic_tokens::SemanticTokenBuilder;
 use crate::state::semantic_tokens::SemanticTokensLegends;
 use crate::state::semantic_tokens::disabled_ranges_for_module;
 use crate::state::state::Transaction;
 
 impl Transaction<'_> {
+    /// Classify an attribute using its value type and effective definition.
+    fn attribute_semantic_token_type(
+        &self,
+        handle: &Handle,
+        attr: &ExprAttribute,
+    ) -> Option<SemanticTokenTypes> {
+        let ty = self.get_type_trace(handle, attr.range())?;
+        let kind = attribute_symbol_kind_from_type(&ty)
+            .to_lsp_semantic_token_type_with_modifiers()
+            .0;
+        let name = &attr.attr.id;
+        // Enum-valued fields can have the same type as enum members. Check
+        // the winning declaration on each receiver branch, not shadowed bases.
+        if ty
+            .as_unions()
+            .iter()
+            .all(|ty| matches!(ty, Type::Literal(lit) if matches!(lit.value, Lit::Enum(_))))
+            && let Some(base) = self.get_type_trace(handle, attr.value.range())
+            && self.ad_hoc_solve(handle, "semantic_token_enum_member", |solver| {
+                solver
+                    .attribute_lookup_classes(&base)
+                    .is_some_and(|classes| {
+                        !classes.is_empty()
+                            && classes.iter().all(|cls| {
+                                solver
+                                    .get_class_member_with_defining_class(cls, name)
+                                    .is_some_and(|member| {
+                                        solver
+                                            .get_enum_member(&member.defining_class, name)
+                                            .is_some()
+                                    })
+                            })
+                    })
+            }) == Some(true)
+        {
+            return Some(SemanticTokenTypes::EnumMember);
+        }
+        Some(kind)
+    }
+
     pub fn semantic_tokens(
         &self,
         handle: &Handle,
@@ -38,7 +84,7 @@ impl Transaction<'_> {
 
         builder.process_ast(
             &ast,
-            &|range| self.get_type_trace(handle, range),
+            &|attr| self.attribute_semantic_token_type(handle, attr),
             &|key: &Key| {
                 let find_preference = FindPreference {
                     import_behavior: ImportBehavior::StopAtRenamedImports,

pyrefly/lib/state/semantic_tokens.rs

@@ -19,8 +19,6 @@ use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_python::symbol_kind::SymbolKind;
 use pyrefly_python::sys_info::SysInfo;
-use pyrefly_types::literal::Lit;
-use pyrefly_types::types::Type;
 use pyrefly_util::visit::Visit as _;
 use regex::Regex;
 use ruff_python_ast::Arguments;
@@ -45,7 +43,6 @@ use ruff_text_size::TextRange;
 use ruff_text_size::TextSize;
 
 use crate::binding::binding::Key;
-use crate::state::lsp::attribute_symbol_kind_from_type;
 
 const SELF_PARAMETER_MODIFIER: SemanticTokenModifiers =
     SemanticTokenModifiers::Custom(Cow::Borrowed("selfParameter"));
@@ -399,35 +396,6 @@ fn collect_interpolated_string_ranges(
     }
 }
 
-/// 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) -> SemanticTokenTypes {
-    match ty {
-        Type::Union(union) => {
-            let mut members = union.members.into_iter();
-            let Some(first) = members.next() else {
-                return SemanticTokenTypes::Property;
-            };
-            let kind = attribute_semantic_token_type(first);
-            if kind == SemanticTokenTypes::Property {
-                return SemanticTokenTypes::Property;
-            }
-            if members.all(|member| attribute_semantic_token_type(member) == kind) {
-                kind
-            } else {
-                SemanticTokenTypes::Property
-            }
-        }
-        Type::Literal(lit) if matches!(lit.value, Lit::Enum(_)) => SemanticTokenTypes::EnumMember,
-        _ => {
-            attribute_symbol_kind_from_type(&ty)
-                .to_lsp_semantic_token_type_with_modifiers()
-                .0
-        }
-    }
-}
-
 pub struct SemanticTokenWithFullRange {
     pub range: TextRange,
     pub token_type: SemanticTokenTypes,
@@ -669,24 +637,10 @@ impl SemanticTokenBuilder {
         });
     }
 
-    fn process_attribute_expr(
-        &mut self,
-        attr: &ExprAttribute,
-        get_type_of_attribute: &dyn Fn(TextRange) -> Option<Type>,
-        get_symbol_kind: &dyn Fn(&Key) -> Option<(ModuleName, SymbolKind)>,
-    ) {
-        let kind = get_type_of_attribute(attr.range())
-            .map(attribute_semantic_token_type)
-            .unwrap_or(SemanticTokenTypes::Property);
-        self.push_if_in_range(attr.attr.range(), kind, Vec::new());
-        attr.value
-            .visit(&mut |x| self.process_expr(x, get_type_of_attribute, get_symbol_kind));
-    }
-
     fn process_expr(
         &mut self,
         x: &Expr,
-        get_type_of_attribute: &dyn Fn(TextRange) -> Option<Type>,
+        get_attribute_kind: &dyn Fn(&ExprAttribute) -> Option<SemanticTokenTypes>,
         get_symbol_kind: &dyn Fn(&Key) -> Option<(ModuleName, SymbolKind)>,
     ) {
         match x {
@@ -714,46 +668,47 @@ impl SemanticTokenBuilder {
             }
             Expr::Call(call) => {
                 self.process_arguments(&call.arguments);
-                x.recurse(&mut |x| self.process_expr(x, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |x| self.process_expr(x, get_attribute_kind, get_symbol_kind));
             }
             Expr::Attribute(attr) => {
-                self.process_attribute_expr(attr, get_type_of_attribute, get_symbol_kind);
+                if range_overlaps(self.limit_range, attr.attr.range()) {
+                    let kind = get_attribute_kind(attr).unwrap_or(SemanticTokenTypes::Property);
+                    self.push_if_in_range(attr.attr.range(), kind, Vec::new());
+                }
+                attr.value
+                    .visit(&mut |x| self.process_expr(x, get_attribute_kind, get_symbol_kind));
             }
             // Comprehensions need special handling because the Visit trait doesn't visit targets
             Expr::ListComp(list_comp) => {
                 for comp in &list_comp.generators {
-                    comp.target.visit(&mut |e| {
-                        self.process_expr(e, get_type_of_attribute, get_symbol_kind)
-                    });
+                    comp.target
+                        .visit(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
                 }
-                x.recurse(&mut |e| self.process_expr(e, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
             }
             Expr::SetComp(set_comp) => {
                 for comp in &set_comp.generators {
-                    comp.target.visit(&mut |e| {
-                        self.process_expr(e, get_type_of_attribute, get_symbol_kind)
-                    });
+                    comp.target
+                        .visit(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
                 }
-                x.recurse(&mut |e| self.process_expr(e, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
             }
             Expr::DictComp(dict_comp) => {
                 for comp in &dict_comp.generators {
-                    comp.target.visit(&mut |e| {
-                        self.process_expr(e, get_type_of_attribute, get_symbol_kind)
-                    });
+                    comp.target
+                        .visit(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
                 }
-                x.recurse(&mut |e| self.process_expr(e, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
             }
             Expr::Generator(generator) => {
                 for comp in &generator.generators {
-                    comp.target.visit(&mut |e| {
-                        self.process_expr(e, get_type_of_attribute, get_symbol_kind)
-                    });
+                    comp.target
+                        .visit(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
                 }
-                x.recurse(&mut |e| self.process_expr(e, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
             }
             _ => {
-                x.recurse(&mut |x| self.process_expr(x, get_type_of_attribute, get_symbol_kind));
+                x.recurse(&mut |x| self.process_expr(x, get_attribute_kind, get_symbol_kind));
             }
         }
     }
@@ -949,13 +904,13 @@ impl SemanticTokenBuilder {
     pub fn process_ast(
         &mut self,
         ast: &ModModule,
-        get_type_of_attribute: &dyn Fn(TextRange) -> Option<Type>,
+        get_attribute_kind: &dyn Fn(&ExprAttribute) -> Option<SemanticTokenTypes>,
         get_symbol_kind: &dyn Fn(&Key) -> Option<(ModuleName, SymbolKind)>,
     ) {
         for s in &ast.body {
             self.process_stmt(s, false, get_symbol_kind);
         }
-        ast.visit(&mut |e| self.process_expr(e, get_type_of_attribute, get_symbol_kind));
+        ast.visit(&mut |e| self.process_expr(e, get_attribute_kind, get_symbol_kind));
     }
 
     pub fn all_tokens_sorted(self) -> Vec<SemanticTokenWithFullRange> {