mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-27 12:29:28 +00:00
Implement visitation of type aliases and parameters (#5927)
<!-- 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
GitHub
parent
3000a47fe8
commit
389fe13c93
18 changed files with 837 additions and 8 deletions
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue