diff --git a/crates/ruff_python_ast/src/node.rs b/crates/ruff_python_ast/src/node.rs index e37467c324..076c299885 100644 --- a/crates/ruff_python_ast/src/node.rs +++ b/crates/ruff_python_ast/src/node.rs @@ -1,3 +1,4 @@ +use ast::{TypeParam, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple}; use ruff_text_size::TextRange; use rustpython_ast::{ Alias, Arg, ArgWithDefault, Arguments, Comprehension, Decorator, ExceptHandler, Keyword, @@ -100,6 +101,9 @@ pub enum AnyNode { MatchCase(MatchCase), Decorator(Decorator), ElifElseClause(ast::ElifElseClause), + TypeParamTypeVar(ast::TypeParamTypeVar), + TypeParamTypeVarTuple(ast::TypeParamTypeVarTuple), + TypeParamParamSpec(ast::TypeParamParamSpec), } impl AnyNode { @@ -184,6 +188,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -269,6 +276,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -354,6 +364,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -439,6 +452,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -524,6 +540,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -609,6 +628,9 @@ impl AnyNode { | AnyNode::WithItem(_) | AnyNode::MatchCase(_) | AnyNode::Decorator(_) + | AnyNode::TypeParamTypeVar(_) + | AnyNode::TypeParamTypeVarTuple(_) + | AnyNode::TypeParamParamSpec(_) | AnyNode::ElifElseClause(_) => None, } } @@ -717,6 +739,9 @@ impl AnyNode { Self::WithItem(node) => AnyNodeRef::WithItem(node), Self::MatchCase(node) => AnyNodeRef::MatchCase(node), Self::Decorator(node) => AnyNodeRef::Decorator(node), + Self::TypeParamTypeVar(node) => AnyNodeRef::TypeParamTypeVar(node), + Self::TypeParamTypeVarTuple(node) => AnyNodeRef::TypeParamTypeVarTuple(node), + Self::TypeParamParamSpec(node) => AnyNodeRef::TypeParamParamSpec(node), Self::ElifElseClause(node) => AnyNodeRef::ElifElseClause(node), } } @@ -2941,7 +2966,90 @@ impl AstNode for Decorator { AnyNode::from(self) } } +impl AstNode for ast::TypeParamTypeVar { + fn cast(kind: AnyNode) -> Option + where + Self: Sized, + { + if let AnyNode::TypeParamTypeVar(node) = kind { + Some(node) + } else { + None + } + } + fn cast_ref(kind: AnyNodeRef) -> Option<&Self> { + if let AnyNodeRef::TypeParamTypeVar(node) = kind { + Some(node) + } else { + None + } + } + + fn as_any_node_ref(&self) -> AnyNodeRef { + AnyNodeRef::from(self) + } + + fn into_any_node(self) -> AnyNode { + AnyNode::from(self) + } +} +impl AstNode for ast::TypeParamTypeVarTuple { + fn cast(kind: AnyNode) -> Option + where + Self: Sized, + { + if let AnyNode::TypeParamTypeVarTuple(node) = kind { + Some(node) + } else { + None + } + } + + fn cast_ref(kind: AnyNodeRef) -> Option<&Self> { + if let AnyNodeRef::TypeParamTypeVarTuple(node) = kind { + Some(node) + } else { + None + } + } + + fn as_any_node_ref(&self) -> AnyNodeRef { + AnyNodeRef::from(self) + } + + fn into_any_node(self) -> AnyNode { + AnyNode::from(self) + } +} +impl AstNode for ast::TypeParamParamSpec { + fn cast(kind: AnyNode) -> Option + where + Self: Sized, + { + if let AnyNode::TypeParamParamSpec(node) = kind { + Some(node) + } else { + None + } + } + + fn cast_ref(kind: AnyNodeRef) -> Option<&Self> { + if let AnyNodeRef::TypeParamParamSpec(node) = kind { + Some(node) + } else { + None + } + } + + fn as_any_node_ref(&self) -> AnyNodeRef { + AnyNodeRef::from(self) + } + + fn into_any_node(self) -> AnyNode { + AnyNode::from(self) + } +} impl From for AnyNode { fn from(stmt: Stmt) -> Self { match stmt { @@ -3518,6 +3626,23 @@ impl From for AnyNode { AnyNode::Decorator(node) } } +impl From for AnyNode { + fn from(node: TypeParamTypeVar) -> Self { + AnyNode::TypeParamTypeVar(node) + } +} + +impl From for AnyNode { + fn from(node: TypeParamTypeVarTuple) -> Self { + AnyNode::TypeParamTypeVarTuple(node) + } +} + +impl From for AnyNode { + fn from(node: TypeParamParamSpec) -> Self { + AnyNode::TypeParamParamSpec(node) + } +} impl Ranged for AnyNode { fn range(&self) -> TextRange { @@ -3600,6 +3725,9 @@ impl Ranged for AnyNode { AnyNode::WithItem(node) => node.range(), AnyNode::MatchCase(node) => node.range(), AnyNode::Decorator(node) => node.range(), + AnyNode::TypeParamTypeVar(node) => node.range(), + AnyNode::TypeParamTypeVarTuple(node) => node.range(), + AnyNode::TypeParamParamSpec(node) => node.range(), AnyNode::ElifElseClause(node) => node.range(), } } @@ -3685,6 +3813,9 @@ pub enum AnyNodeRef<'a> { WithItem(&'a WithItem), MatchCase(&'a MatchCase), Decorator(&'a Decorator), + TypeParamTypeVar(&'a ast::TypeParamTypeVar), + TypeParamTypeVarTuple(&'a ast::TypeParamTypeVarTuple), + TypeParamParamSpec(&'a ast::TypeParamParamSpec), ElifElseClause(&'a ast::ElifElseClause), } @@ -3769,6 +3900,9 @@ impl AnyNodeRef<'_> { AnyNodeRef::WithItem(node) => NonNull::from(*node).cast(), AnyNodeRef::MatchCase(node) => NonNull::from(*node).cast(), AnyNodeRef::Decorator(node) => NonNull::from(*node).cast(), + AnyNodeRef::TypeParamTypeVar(node) => NonNull::from(*node).cast(), + AnyNodeRef::TypeParamTypeVarTuple(node) => NonNull::from(*node).cast(), + AnyNodeRef::TypeParamParamSpec(node) => NonNull::from(*node).cast(), AnyNodeRef::ElifElseClause(node) => NonNull::from(*node).cast(), } } @@ -3859,6 +3993,9 @@ impl AnyNodeRef<'_> { AnyNodeRef::WithItem(_) => NodeKind::WithItem, AnyNodeRef::MatchCase(_) => NodeKind::MatchCase, AnyNodeRef::Decorator(_) => NodeKind::Decorator, + AnyNodeRef::TypeParamTypeVar(_) => NodeKind::TypeParamTypeVar, + AnyNodeRef::TypeParamTypeVarTuple(_) => NodeKind::TypeParamTypeVarTuple, + AnyNodeRef::TypeParamParamSpec(_) => NodeKind::TypeParamParamSpec, AnyNodeRef::ElifElseClause(_) => NodeKind::ElifElseClause, } } @@ -3944,6 +4081,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4029,6 +4169,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4114,6 +4257,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4199,6 +4345,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4284,6 +4433,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4369,6 +4521,9 @@ impl AnyNodeRef<'_> { | AnyNodeRef::WithItem(_) | AnyNodeRef::MatchCase(_) | AnyNodeRef::Decorator(_) + | AnyNodeRef::TypeParamTypeVar(_) + | AnyNodeRef::TypeParamTypeVarTuple(_) + | AnyNodeRef::TypeParamParamSpec(_) | AnyNodeRef::ElifElseClause(_) => false, } } @@ -4828,6 +4983,24 @@ impl<'a> From<&'a Decorator> for AnyNodeRef<'a> { } } +impl<'a> From<&'a TypeParamTypeVar> for AnyNodeRef<'a> { + fn from(node: &'a TypeParamTypeVar) -> Self { + AnyNodeRef::TypeParamTypeVar(node) + } +} + +impl<'a> From<&'a TypeParamTypeVarTuple> for AnyNodeRef<'a> { + fn from(node: &'a TypeParamTypeVarTuple) -> Self { + AnyNodeRef::TypeParamTypeVarTuple(node) + } +} + +impl<'a> From<&'a TypeParamParamSpec> for AnyNodeRef<'a> { + fn from(node: &'a TypeParamParamSpec) -> Self { + AnyNodeRef::TypeParamParamSpec(node) + } +} + impl<'a> From<&'a Stmt> for AnyNodeRef<'a> { fn from(stmt: &'a Stmt) -> Self { match stmt { @@ -4923,6 +5096,16 @@ impl<'a> From<&'a Pattern> for AnyNodeRef<'a> { } } +impl<'a> From<&'a TypeParam> for AnyNodeRef<'a> { + fn from(type_param: &'a TypeParam) -> Self { + match type_param { + TypeParam::TypeVar(node) => AnyNodeRef::TypeParamTypeVar(node), + TypeParam::TypeVarTuple(node) => AnyNodeRef::TypeParamTypeVarTuple(node), + TypeParam::ParamSpec(node) => AnyNodeRef::TypeParamParamSpec(node), + } + } +} + impl<'a> From<&'a ExceptHandler> for AnyNodeRef<'a> { fn from(handler: &'a ExceptHandler) -> Self { match handler { @@ -5064,6 +5247,9 @@ impl Ranged for AnyNodeRef<'_> { AnyNodeRef::MatchCase(node) => node.range(), AnyNodeRef::Decorator(node) => node.range(), AnyNodeRef::ElifElseClause(node) => node.range(), + AnyNodeRef::TypeParamTypeVar(node) => node.range(), + AnyNodeRef::TypeParamTypeVarTuple(node) => node.range(), + AnyNodeRef::TypeParamParamSpec(node) => node.range(), } } } @@ -5149,4 +5335,7 @@ pub enum NodeKind { MatchCase, Decorator, ElifElseClause, + TypeParamTypeVar, + TypeParamTypeVarTuple, + TypeParamParamSpec, } diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__class_type_parameters.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__class_type_parameters.snap new file mode 100644 index 0000000000..929aa88b29 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__class_type_parameters.snap @@ -0,0 +1,13 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtClassDef + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - StmtExpr + - ExprConstant + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__compare.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__compare.snap new file mode 100644 index 0000000000..ce293eef68 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__compare.snap @@ -0,0 +1,12 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtExpr + - ExprCompare + - ExprConstant + - Lt + - Lt + - ExprName + - ExprConstant + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__decorators.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__decorators.snap new file mode 100644 index 0000000000..a0f99c2447 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__decorators.snap @@ -0,0 +1,12 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtFunctionDef + - ExprName + - Arguments + - StmtPass +- StmtClassDef + - ExprName + - StmtPass + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__dict_comprehension.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__dict_comprehension.snap new file mode 100644 index 0000000000..baf55d5a9a --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__dict_comprehension.snap @@ -0,0 +1,15 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtExpr + - ExprDictComp + - Comprehension + - ExprName + - ExprName + - ExprName + - ExprBinOp + - ExprName + - Pow + - ExprConstant + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_arguments.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_arguments.snap new file mode 100644 index 0000000000..88d08c74e4 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_arguments.snap @@ -0,0 +1,19 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtFunctionDef + - Arguments + - ExprConstant + - ExprConstant + - ExprConstant + - Arg + - Arg + - Arg + - Arg + - Arg + - Arg + - Arg + - Arg + - StmtPass + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_positional_only_with_default.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_positional_only_with_default.snap new file mode 100644 index 0000000000..d77750d81a --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_positional_only_with_default.snap @@ -0,0 +1,14 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtFunctionDef + - Arguments + - ExprConstant + - ExprConstant + - Arg + - Arg + - Arg + - Arg + - StmtPass + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_type_parameters.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_type_parameters.snap new file mode 100644 index 0000000000..290000d4d3 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_type_parameters.snap @@ -0,0 +1,14 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtFunctionDef + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - Arguments + - StmtExpr + - ExprConstant + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__list_comprehension.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__list_comprehension.snap new file mode 100644 index 0000000000..464f2cf03f --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__list_comprehension.snap @@ -0,0 +1,11 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtExpr + - ExprListComp + - Comprehension + - ExprName + - ExprName + - ExprName + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__match_class_pattern.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__match_class_pattern.snap new file mode 100644 index 0000000000..ce3f84813b --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__match_class_pattern.snap @@ -0,0 +1,27 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtMatch + - ExprName + - MatchCase + - PatternMatchClass + - ExprName + - PatternMatchValue + - ExprConstant + - PatternMatchValue + - ExprConstant + - StmtExpr + - ExprConstant + - MatchCase + - PatternMatchClass + - ExprName + - PatternMatchValue + - ExprConstant + - PatternMatchValue + - ExprConstant + - PatternMatchValue + - ExprConstant + - StmtExpr + - ExprConstant + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__set_comprehension.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__set_comprehension.snap new file mode 100644 index 0000000000..716eb2f350 --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__set_comprehension.snap @@ -0,0 +1,11 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtExpr + - ExprSetComp + - Comprehension + - ExprName + - ExprName + - ExprName + diff --git a/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__type_aliases.snap b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__type_aliases.snap new file mode 100644 index 0000000000..6c0675c64c --- /dev/null +++ b/crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__type_aliases.snap @@ -0,0 +1,15 @@ +--- +source: crates/ruff_python_ast/src/visitor.rs +expression: trace +--- +- StmtTypeAlias + - ExprSubscript + - ExprName + - ExprName + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - ExprName + diff --git a/crates/ruff_python_ast/src/visitor.rs b/crates/ruff_python_ast/src/visitor.rs index bf2f87ac92..8fb17abbd1 100644 --- a/crates/ruff_python_ast/src/visitor.rs +++ b/crates/ruff_python_ast/src/visitor.rs @@ -5,7 +5,8 @@ pub mod preorder; use rustpython_ast::ElifElseClause; use rustpython_parser::ast::{ self, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, Decorator, ExceptHandler, Expr, - ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, UnaryOp, WithItem, + ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, TypeParam, TypeParamTypeVar, UnaryOp, + WithItem, }; /// A trait for AST visitors. Visits all nodes in the AST recursively in evaluation-order. @@ -67,6 +68,9 @@ pub trait Visitor<'a> { fn visit_with_item(&mut self, with_item: &'a WithItem) { walk_with_item(self, with_item); } + fn visit_type_param(&mut self, type_param: &'a TypeParam) { + walk_type_param(self, type_param); + } fn visit_match_case(&mut self, match_case: &'a MatchCase) { walk_match_case(self, match_case); } @@ -104,11 +108,15 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) { body, decorator_list, returns, + type_params, .. }) => { for decorator in decorator_list { visitor.visit_decorator(decorator); } + for type_param in type_params { + visitor.visit_type_param(type_param); + } visitor.visit_arguments(args); for expr in returns { visitor.visit_annotation(expr); @@ -120,11 +128,15 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) { body, decorator_list, returns, + type_params, .. }) => { for decorator in decorator_list { visitor.visit_decorator(decorator); } + for type_param in type_params { + visitor.visit_type_param(type_param); + } visitor.visit_arguments(args); for expr in returns { visitor.visit_annotation(expr); @@ -136,11 +148,15 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) { keywords, body, decorator_list, + type_params, .. }) => { for decorator in decorator_list { visitor.visit_decorator(decorator); } + for type_param in type_params { + visitor.visit_type_param(type_param); + } for expr in bases { visitor.visit_expr(expr); } @@ -165,6 +181,18 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) { visitor.visit_expr(expr); } } + Stmt::TypeAlias(ast::StmtTypeAlias { + range: _range, + name, + type_params, + value, + }) => { + visitor.visit_expr(value); + for type_param in type_params { + visitor.visit_type_param(type_param); + } + visitor.visit_expr(name); + } Stmt::Assign(ast::StmtAssign { targets, value, .. }) => { visitor.visit_expr(value); for expr in targets { @@ -334,7 +362,6 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) { range: _range, }) => visitor.visit_expr(value), Stmt::Pass(_) | Stmt::Break(_) | Stmt::Continue(_) => {} - Stmt::TypeAlias(_) => todo!(), } } @@ -670,6 +697,21 @@ pub fn walk_with_item<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, with_item: & } } +pub fn walk_type_param<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, type_param: &'a TypeParam) { + match type_param { + TypeParam::TypeVar(TypeParamTypeVar { + bound, + name: _, + range: _, + }) => { + if let Some(expr) = bound { + visitor.visit_expr(expr); + } + } + TypeParam::TypeVarTuple(_) | TypeParam::ParamSpec(_) => {} + } +} + pub fn walk_match_case<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, match_case: &'a MatchCase) { visitor.visit_pattern(&match_case.pattern); if let Some(expr) = &match_case.guard { @@ -754,3 +796,305 @@ pub fn walk_cmp_op<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, cmp_op: &'a Cmp #[allow(unused_variables)] pub fn walk_alias<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, alias: &'a Alias) {} + +#[cfg(test)] +mod tests { + use std::fmt::{Debug, Write}; + + use insta::assert_snapshot; + use rustpython_parser::lexer::lex; + use rustpython_parser::{ast, parse_tokens, Mode}; + + use crate::node::AnyNodeRef; + use crate::visitor::{ + walk_alias, walk_arg, walk_arguments, walk_comprehension, walk_except_handler, walk_expr, + walk_keyword, walk_match_case, walk_pattern, walk_stmt, walk_type_param, walk_with_item, + Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, ExceptHandler, Expr, Keyword, + MatchCase, Operator, Pattern, Stmt, TypeParam, UnaryOp, Visitor, WithItem, + }; + + #[test] + fn function_arguments() { + let source = r#"def a(b, c,/, d, e = 20, *args, named=5, other=20, **kwargs): pass"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn function_positional_only_with_default() { + let source = r#"def a(b, c = 34,/, e = 20, *args): pass"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn compare() { + let source = r#"4 < x < 5"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn list_comprehension() { + let source = "[x for x in numbers]"; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn dict_comprehension() { + let source = "{x: x**2 for x in numbers}"; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn set_comprehension() { + let source = "{x for x in numbers}"; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn match_class_pattern() { + let source = r#" +match x: + case Point2D(0, 0): + ... + case Point3D(x=0, y=0, z=0): + ... +"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn decorators() { + let source = r#" +@decorator +def a(): + pass + +@test +class A: + pass +"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn type_aliases() { + let source = r#"type X[T: str, U, *Ts, **P] = list[T]"#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn class_type_parameters() { + let source = r#"class X[T: str, U, *Ts, **P]: ..."#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn function_type_parameters() { + let source = r#"def X[T: str, U, *Ts, **P](): ..."#; + + let trace = trace_visitation(source); + + assert_snapshot!(trace); + } + + fn trace_visitation(source: &str) -> String { + let tokens = lex(source, Mode::Module); + let parsed = parse_tokens(tokens, Mode::Module, "test.py").unwrap(); + + let mut visitor = RecordVisitor::default(); + walk_module(&mut visitor, &parsed); + + visitor.output + } + + fn walk_module<'a, V>(visitor: &mut V, module: &'a ast::Mod) + where + V: Visitor<'a> + ?Sized, + { + match module { + ast::Mod::Module(ast::ModModule { + body, + range: _, + type_ignores: _, + }) => { + visitor.visit_body(body); + } + ast::Mod::Interactive(ast::ModInteractive { body, range: _ }) => { + visitor.visit_body(body); + } + ast::Mod::Expression(ast::ModExpression { body, range: _ }) => visitor.visit_expr(body), + ast::Mod::FunctionType(ast::ModFunctionType { + range: _, + argtypes, + returns, + }) => { + for arg_type in argtypes { + visitor.visit_expr(arg_type); + } + + visitor.visit_expr(returns); + } + } + } + + /// Emits a `tree` with a node for every visited AST node (labelled by the AST node's kind) + /// and leafs for attributes. + #[derive(Default)] + struct RecordVisitor { + depth: usize, + output: String, + } + + impl RecordVisitor { + fn enter_node<'a, T>(&mut self, node: T) + where + T: Into>, + { + self.emit(&node.into().kind()); + self.depth += 1; + } + + fn exit_node(&mut self) { + self.depth -= 1; + } + + fn emit(&mut self, text: &dyn Debug) { + for _ in 0..self.depth { + self.output.push_str(" "); + } + + writeln!(self.output, "- {text:?}").unwrap(); + } + } + + impl Visitor<'_> for RecordVisitor { + fn visit_stmt(&mut self, stmt: &Stmt) { + self.enter_node(stmt); + walk_stmt(self, stmt); + self.exit_node(); + } + + fn visit_annotation(&mut self, expr: &Expr) { + self.enter_node(expr); + walk_expr(self, expr); + self.exit_node(); + } + + fn visit_expr(&mut self, expr: &Expr) { + self.enter_node(expr); + walk_expr(self, expr); + self.exit_node(); + } + + fn visit_bool_op(&mut self, bool_op: &BoolOp) { + self.emit(&bool_op); + } + + fn visit_operator(&mut self, operator: &Operator) { + self.emit(&operator); + } + + fn visit_unary_op(&mut self, unary_op: &UnaryOp) { + self.emit(&unary_op); + } + + fn visit_cmp_op(&mut self, cmp_op: &CmpOp) { + self.emit(&cmp_op); + } + + fn visit_comprehension(&mut self, comprehension: &Comprehension) { + self.enter_node(comprehension); + walk_comprehension(self, comprehension); + self.exit_node(); + } + + fn visit_except_handler(&mut self, except_handler: &ExceptHandler) { + self.enter_node(except_handler); + walk_except_handler(self, except_handler); + self.exit_node(); + } + + fn visit_format_spec(&mut self, format_spec: &Expr) { + self.enter_node(format_spec); + walk_expr(self, format_spec); + self.exit_node(); + } + + fn visit_arguments(&mut self, arguments: &Arguments) { + self.enter_node(arguments); + walk_arguments(self, arguments); + self.exit_node(); + } + + fn visit_arg(&mut self, arg: &Arg) { + self.enter_node(arg); + walk_arg(self, arg); + self.exit_node(); + } + + fn visit_keyword(&mut self, keyword: &Keyword) { + self.enter_node(keyword); + walk_keyword(self, keyword); + self.exit_node(); + } + + fn visit_alias(&mut self, alias: &Alias) { + self.enter_node(alias); + walk_alias(self, alias); + self.exit_node(); + } + + fn visit_with_item(&mut self, with_item: &WithItem) { + self.enter_node(with_item); + walk_with_item(self, with_item); + self.exit_node(); + } + + fn visit_match_case(&mut self, match_case: &MatchCase) { + self.enter_node(match_case); + walk_match_case(self, match_case); + self.exit_node(); + } + + fn visit_pattern(&mut self, pattern: &Pattern) { + self.enter_node(pattern); + walk_pattern(self, pattern); + self.exit_node(); + } + + fn visit_type_param(&mut self, type_param: &TypeParam) { + self.enter_node(type_param); + walk_type_param(self, type_param); + self.exit_node(); + } + } +} diff --git a/crates/ruff_python_ast/src/visitor/preorder.rs b/crates/ruff_python_ast/src/visitor/preorder.rs index ae9209d21f..b3570b205c 100644 --- a/crates/ruff_python_ast/src/visitor/preorder.rs +++ b/crates/ruff_python_ast/src/visitor/preorder.rs @@ -1,7 +1,8 @@ use rustpython_ast::{ArgWithDefault, ElifElseClause, Mod, TypeIgnore}; use rustpython_parser::ast::{ self, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, Constant, Decorator, ExceptHandler, - Expr, Keyword, MatchCase, Operator, Pattern, Stmt, UnaryOp, WithItem, + Expr, Keyword, MatchCase, Operator, Pattern, Stmt, TypeParam, TypeParamTypeVar, UnaryOp, + WithItem, }; /// Visitor that traverses all nodes recursively in pre-order. @@ -80,6 +81,10 @@ pub trait PreorderVisitor<'a> { walk_with_item(self, with_item); } + fn visit_type_param(&mut self, type_param: &'a TypeParam) { + walk_type_param(self, type_param); + } + fn visit_match_case(&mut self, match_case: &'a MatchCase) { walk_match_case(self, match_case); } @@ -156,6 +161,7 @@ where body, decorator_list, returns, + type_params, .. }) | Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { @@ -163,12 +169,17 @@ where body, decorator_list, returns, + type_params, .. }) => { for decorator in decorator_list { visitor.visit_decorator(decorator); } + for type_param in type_params { + visitor.visit_type_param(type_param); + } + visitor.visit_arguments(args); for expr in returns { @@ -183,12 +194,17 @@ where keywords, body, decorator_list, + type_params, .. }) => { for decorator in decorator_list { visitor.visit_decorator(decorator); } + for type_param in type_params { + visitor.visit_type_param(type_param); + } + for expr in bases { visitor.visit_expr(expr); } @@ -218,6 +234,19 @@ where } } + Stmt::TypeAlias(ast::StmtTypeAlias { + range: _range, + name, + type_params, + value, + }) => { + visitor.visit_expr(name); + for type_param in type_params { + visitor.visit_type_param(type_param); + } + visitor.visit_expr(value); + } + Stmt::Assign(ast::StmtAssign { targets, value, @@ -400,7 +429,6 @@ where | Stmt::Continue(_) | Stmt::Global(_) | Stmt::Nonlocal(_) => {} - Stmt::TypeAlias(_) => todo!(), } } @@ -805,6 +833,24 @@ where } } +pub fn walk_type_param<'a, V>(visitor: &mut V, type_param: &'a TypeParam) +where + V: PreorderVisitor<'a> + ?Sized, +{ + match type_param { + TypeParam::TypeVar(TypeParamTypeVar { + bound, + name: _, + range: _, + }) => { + if let Some(expr) = bound { + visitor.visit_expr(expr); + } + } + TypeParam::TypeVarTuple(_) | TypeParam::ParamSpec(_) => {} + } +} + pub fn walk_match_case<'a, V>(visitor: &mut V, match_case: &'a MatchCase) where V: PreorderVisitor<'a> + ?Sized, @@ -947,9 +993,9 @@ mod tests { use crate::visitor::preorder::{ walk_alias, walk_arg, walk_arguments, walk_comprehension, walk_except_handler, walk_expr, walk_keyword, walk_match_case, walk_module, walk_pattern, walk_stmt, walk_type_ignore, - walk_with_item, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, Constant, - ExceptHandler, Expr, Keyword, MatchCase, Mod, Operator, Pattern, PreorderVisitor, Stmt, - TypeIgnore, UnaryOp, WithItem, + walk_type_param, walk_with_item, Alias, Arg, Arguments, BoolOp, CmpOp, Comprehension, + Constant, ExceptHandler, Expr, Keyword, MatchCase, Mod, Operator, Pattern, PreorderVisitor, + Stmt, TypeIgnore, TypeParam, UnaryOp, WithItem, }; #[test] @@ -1038,6 +1084,33 @@ class A: assert_snapshot!(trace); } + #[test] + fn type_aliases() { + let source = r#"type X[T: str, U, *Ts, **P] = list[T]"#; + + let trace = trace_preorder_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn class_type_parameters() { + let source = r#"class X[T: str, U, *Ts, **P]: ..."#; + + let trace = trace_preorder_visitation(source); + + assert_snapshot!(trace); + } + + #[test] + fn function_type_parameters() { + let source = r#"def X[T: str, U, *Ts, **P](): ..."#; + + let trace = trace_preorder_visitation(source); + + assert_snapshot!(trace); + } + fn trace_preorder_visitation(source: &str) -> String { let tokens = lex(source, Mode::Module); let parsed = parse_tokens(tokens, Mode::Module, "test.py").unwrap(); @@ -1188,5 +1261,11 @@ class A: walk_type_ignore(self, type_ignore); self.exit_node(); } + + fn visit_type_param(&mut self, type_param: &TypeParam) { + self.enter_node(type_param); + walk_type_param(self, type_param); + self.exit_node(); + } } } diff --git a/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__class_type_parameters.snap b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__class_type_parameters.snap new file mode 100644 index 0000000000..01cc10719b --- /dev/null +++ b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__class_type_parameters.snap @@ -0,0 +1,15 @@ +--- +source: crates/ruff_python_ast/src/visitor/preorder.rs +expression: trace +--- +- ModModule + - StmtClassDef + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - StmtExpr + - ExprConstant + - Ellipsis + diff --git a/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__function_type_parameters.snap b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__function_type_parameters.snap new file mode 100644 index 0000000000..9c9f0cbd28 --- /dev/null +++ b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__function_type_parameters.snap @@ -0,0 +1,16 @@ +--- +source: crates/ruff_python_ast/src/visitor/preorder.rs +expression: trace +--- +- ModModule + - StmtFunctionDef + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - Arguments + - StmtExpr + - ExprConstant + - Ellipsis + diff --git a/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__type_aliases.snap b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__type_aliases.snap new file mode 100644 index 0000000000..007ee12733 --- /dev/null +++ b/crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__type_aliases.snap @@ -0,0 +1,16 @@ +--- +source: crates/ruff_python_ast/src/visitor/preorder.rs +expression: trace +--- +- ModModule + - StmtTypeAlias + - ExprName + - TypeParamTypeVar + - ExprName + - TypeParamTypeVar + - TypeParamTypeVarTuple + - TypeParamParamSpec + - ExprSubscript + - ExprName + - ExprName + diff --git a/crates/ruff_python_formatter/src/comments/visitor.rs b/crates/ruff_python_formatter/src/comments/visitor.rs index d8131a78fa..87d6b317f6 100644 --- a/crates/ruff_python_formatter/src/comments/visitor.rs +++ b/crates/ruff_python_formatter/src/comments/visitor.rs @@ -3,7 +3,7 @@ use std::iter::Peekable; use ruff_text_size::{TextRange, TextSize}; use rustpython_parser::ast::{ Alias, Arg, ArgWithDefault, Arguments, Comprehension, Decorator, ElifElseClause, ExceptHandler, - Expr, Keyword, MatchCase, Mod, Pattern, Ranged, Stmt, WithItem, + Expr, Keyword, MatchCase, Mod, Pattern, Ranged, Stmt, TypeParam, WithItem, }; use ruff_formatter::{SourceCode, SourceCodeSlice}; @@ -291,6 +291,13 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> { } self.finish_node(elif_else_clause); } + + fn visit_type_param(&mut self, type_param: &'ast TypeParam) { + if self.start_node(type_param).is_traverse() { + walk_type_param(self, type_param); + } + self.finish_node(type_param); + } } fn text_position(comment_range: TextRange, source_code: SourceCode) -> CommentLinePosition {