mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-07-24 13:33:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Implement visitation of type aliases and parameters (#5927)

Browse source 
<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

Part of #5062 
Requires https://github.com/astral-sh/RustPython-Parser/pull/32

Adds visitation of type alias statements and type parameters in class
and function definitions.

Duplicates tests for `PreorderVisitor` into `Visitor` with new
snapshots. Testing required node implementations for the `TypeParam`
enum, which is a chunk of the diff and the reason we need `Ranged`
implementations in
https://github.com/astral-sh/RustPython-Parser/pull/32.

## Test Plan

<!-- How was it tested? -->

Adds unit tests with snapshots.
This commit is contained in:
Zanie Blue 2023-07-25 12:11:26 -05:00 • committed by GitHub
parent 3000a47fe8
commit 389fe13c93
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 837 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

189
crates/ruff_python_ast/src/node.rs
View file

@ -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<Self>
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<Self>
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<Self>
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<Stmt> for AnyNode {
fn from(stmt: Stmt) -> Self {
match stmt {
@ -3518,6 +3626,23 @@ impl From<Decorator> for AnyNode {
AnyNode::Decorator(node)
}
}
impl From<TypeParamTypeVar> for AnyNode {
fn from(node: TypeParamTypeVar) -> Self {
AnyNode::TypeParamTypeVar(node)
}
}
impl From<TypeParamTypeVarTuple> for AnyNode {
fn from(node: TypeParamTypeVarTuple) -> Self {
AnyNode::TypeParamTypeVarTuple(node)
}
}
impl From<TypeParamParamSpec> 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,
}

13
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__class_type_parameters.snap Normal file
View file

@ -0,0 +1,13 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtClassDef
- TypeParamTypeVar
- ExprName
- TypeParamTypeVar
- TypeParamTypeVarTuple
- TypeParamParamSpec
- StmtExpr
- ExprConstant

12
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__compare.snap Normal file
View file

@ -0,0 +1,12 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtExpr
- ExprCompare
- ExprConstant
- Lt
- Lt
- ExprName
- ExprConstant

12
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__decorators.snap Normal file
View file

@ -0,0 +1,12 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtFunctionDef
- ExprName
- Arguments
- StmtPass
- StmtClassDef
- ExprName
- StmtPass

15
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__dict_comprehension.snap Normal file
View file

@ -0,0 +1,15 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtExpr
- ExprDictComp
- Comprehension
- ExprName
- ExprName
- ExprName
- ExprBinOp
- ExprName
- Pow
- ExprConstant

19
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_arguments.snap Normal file
View file

@ -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

14
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_positional_only_with_default.snap Normal file
View file

@ -0,0 +1,14 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtFunctionDef
- Arguments
- ExprConstant
- ExprConstant
- Arg
- Arg
- Arg
- Arg
- StmtPass

14
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__function_type_parameters.snap Normal file
View file

@ -0,0 +1,14 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtFunctionDef
- TypeParamTypeVar
- ExprName
- TypeParamTypeVar
- TypeParamTypeVarTuple
- TypeParamParamSpec
- Arguments
- StmtExpr
- ExprConstant

11
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__list_comprehension.snap Normal file
View file

@ -0,0 +1,11 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtExpr
- ExprListComp
- Comprehension
- ExprName
- ExprName
- ExprName

27
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__match_class_pattern.snap Normal file
View file

@ -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

11
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__set_comprehension.snap Normal file
View file

@ -0,0 +1,11 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtExpr
- ExprSetComp
- Comprehension
- ExprName
- ExprName
- ExprName

15
crates/ruff_python_ast/src/snapshots/ruff_python_ast__visitor__tests__type_aliases.snap Normal file
View file

@ -0,0 +1,15 @@
---
source: crates/ruff_python_ast/src/visitor.rs
expression: trace
---
- StmtTypeAlias
- ExprSubscript
- ExprName
- ExprName
- TypeParamTypeVar
- ExprName
- TypeParamTypeVar
- TypeParamTypeVarTuple
- TypeParamParamSpec
- ExprName

348
crates/ruff_python_ast/src/visitor.rs
View file

@ -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<AnyNodeRef<'a>>,
{
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();
}
}
}

89
crates/ruff_python_ast/src/visitor/preorder.rs
View file

@ -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();
}
}
}

15
crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__class_type_parameters.snap Normal file
View file

@ -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

16
crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__function_type_parameters.snap Normal file
View file

@ -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

16
crates/ruff_python_ast/src/visitor/snapshots/ruff_python_ast__visitor__preorder__tests__type_aliases.snap Normal file
View file

@ -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

9
crates/ruff_python_formatter/src/comments/visitor.rs
View file

@ -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 {