diff --git a/crates/ty_ide/src/semantic_tokens.rs b/crates/ty_ide/src/semantic_tokens.rs index 2b05548b6e..0fc12d906b 100644 --- a/crates/ty_ide/src/semantic_tokens.rs +++ b/crates/ty_ide/src/semantic_tokens.rs @@ -37,13 +37,12 @@ use bitflags::bitflags; use itertools::Itertools; use ruff_db::files::File; use ruff_db::parsed::parsed_module; -use ruff_python_ast as ast; use ruff_python_ast::visitor::source_order::{ - SourceOrderVisitor, TraversalSignal, walk_expr, walk_stmt, + SourceOrderVisitor, TraversalSignal, walk_arguments, walk_expr, walk_stmt, }; use ruff_python_ast::{ - AnyNodeRef, BytesLiteral, Expr, FString, InterpolatedStringElement, Stmt, StringLiteral, - TypeParam, + self as ast, AnyNodeRef, BytesLiteral, Expr, FString, InterpolatedStringElement, Stmt, + StringLiteral, TypeParam, }; use ruff_text_size::{Ranged, TextLen, TextRange, TextSize}; use std::ops::Deref; @@ -668,14 +667,7 @@ impl SourceOrderVisitor<'_> for SemanticTokenVisitor<'_> { // Handle base classes and type annotations in inheritance if let Some(arguments) = &class.arguments { - // Visit base class arguments - for arg in &arguments.args { - self.visit_expr(arg); - } - // Visit keyword arguments (for metaclass, etc.) - for keyword in &arguments.keywords { - self.visit_expr(&keyword.value); - } + walk_arguments(self, arguments); } let prev_in_class = self.in_class_scope; @@ -1138,6 +1130,22 @@ mod tests { "###); } + #[test] + fn test_semantic_tokens_class_args() { + // This used to cause a panic because of an incorrect + // insertion-order when visiting arguments inside + // class definitions. + let test = SemanticTokenTest::new("class Foo(m=x, m)"); + + let tokens = test.highlight_file(); + + assert_snapshot!(test.to_snapshot(&tokens), @r###" + "Foo" @ 6..9: Class [definition] + "x" @ 12..13: Variable + "m" @ 15..16: Variable + "###); + } + #[test] fn test_semantic_tokens_variables() { let test = SemanticTokenTest::new(