Refactor range from Attributed to Nodes (#4422)

This commit is contained in:
Micha Reiser 2023-05-16 08:36:32 +02:00 committed by GitHub
parent 140e0acf54
commit fa26860296
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
330 changed files with 4816 additions and 3946 deletions

View file

@ -1,5 +1,5 @@
use bitflags::bitflags;
use rustpython_parser::ast::{self, Constant, Expr, ExprKind, Stmt, StmtKind};
use rustpython_parser::ast::{self, Constant, Expr, Stmt};
bitflags! {
#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
@ -18,10 +18,10 @@ where
{
fn add_to_names<'a>(elts: &'a [Expr], names: &mut Vec<&'a str>, flags: &mut AllNamesFlags) {
for elt in elts {
if let ExprKind::Constant(ast::ExprConstant {
if let Expr::Constant(ast::ExprConstant {
value: Constant::Str(value),
..
}) = &elt.node
}) = elt
{
names.push(value);
} else {
@ -30,44 +30,45 @@ where
}
}
fn extract_elts<F>(expr: &Expr, is_builtin: F) -> (Option<&Vec<Expr>>, AllNamesFlags)
fn extract_elts<F>(expr: &Expr, is_builtin: F) -> (Option<&[Expr]>, AllNamesFlags)
where
F: Fn(&str) -> bool,
{
match &expr.node {
ExprKind::List(ast::ExprList { elts, .. }) => {
match expr {
Expr::List(ast::ExprList { elts, .. }) => {
return (Some(elts), AllNamesFlags::empty());
}
ExprKind::Tuple(ast::ExprTuple { elts, .. }) => {
Expr::Tuple(ast::ExprTuple { elts, .. }) => {
return (Some(elts), AllNamesFlags::empty());
}
ExprKind::ListComp(_) => {
Expr::ListComp(_) => {
// Allow comprehensions, even though we can't statically analyze them.
return (None, AllNamesFlags::empty());
}
ExprKind::Call(ast::ExprCall {
Expr::Call(ast::ExprCall {
func,
args,
keywords,
range: _range,
}) => {
// Allow `tuple()` and `list()` calls.
if keywords.is_empty() && args.len() <= 1 {
if let ExprKind::Name(ast::ExprName { id, .. }) = &func.node {
if let Expr::Name(ast::ExprName { id, .. }) = &**func {
let id = id.as_str();
if id == "tuple" || id == "list" {
if is_builtin(id) {
if args.is_empty() {
return (None, AllNamesFlags::empty());
}
match &args[0].node {
ExprKind::List(ast::ExprList { elts, .. })
| ExprKind::Set(ast::ExprSet { elts })
| ExprKind::Tuple(ast::ExprTuple { elts, .. }) => {
match &args[0] {
Expr::List(ast::ExprList { elts, .. })
| Expr::Set(ast::ExprSet { elts, .. })
| Expr::Tuple(ast::ExprTuple { elts, .. }) => {
return (Some(elts), AllNamesFlags::empty());
}
ExprKind::ListComp(_)
| ExprKind::SetComp(_)
| ExprKind::GeneratorExp(_) => {
Expr::ListComp(_)
| Expr::SetComp(_)
| Expr::GeneratorExp(_) => {
// Allow comprehensions, even though we can't statically analyze
// them.
return (None, AllNamesFlags::empty());
@ -87,13 +88,13 @@ where
let mut names: Vec<&str> = vec![];
let mut flags = AllNamesFlags::empty();
if let Some(value) = match &stmt.node {
StmtKind::Assign(ast::StmtAssign { value, .. }) => Some(value),
StmtKind::AnnAssign(ast::StmtAnnAssign { value, .. }) => value.as_ref(),
StmtKind::AugAssign(ast::StmtAugAssign { value, .. }) => Some(value),
if let Some(value) = match stmt {
Stmt::Assign(ast::StmtAssign { value, .. }) => Some(value),
Stmt::AnnAssign(ast::StmtAnnAssign { value, .. }) => value.as_ref(),
Stmt::AugAssign(ast::StmtAugAssign { value, .. }) => Some(value),
_ => None,
} {
if let ExprKind::BinOp(ast::ExprBinOp { left, right, .. }) = &value.node {
if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = &**value {
let mut current_left = left;
let mut current_right = right;
loop {
@ -106,7 +107,7 @@ where
// Process the left side, which can be a "real" value or the "rest" of the
// binary operation.
if let ExprKind::BinOp(ast::ExprBinOp { left, right, .. }) = &current_left.node {
if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = &**current_left {
current_left = left;
current_right = right;
} else {

View file

@ -1,12 +1,12 @@
use rustpython_parser::ast::{self, Expr, ExprKind};
use rustpython_parser::ast::{self, Expr};
use smallvec::smallvec;
/// A representation of a qualified name, like `typing.List`.
pub type CallPath<'a> = smallvec::SmallVec<[&'a str; 8]>;
fn collect_call_path_inner<'a>(expr: &'a Expr, parts: &mut CallPath<'a>) -> bool {
match &expr.node {
ExprKind::Attribute(ast::ExprAttribute { value, attr, .. }) => {
match expr {
Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
if collect_call_path_inner(value, parts) {
parts.push(attr.as_str());
true
@ -14,7 +14,7 @@ fn collect_call_path_inner<'a>(expr: &'a Expr, parts: &mut CallPath<'a>) -> bool
false
}
}
ExprKind::Name(ast::ExprName { id, .. }) => {
Expr::Name(ast::ExprName { id, .. }) => {
parts.push(id.as_str());
true
}

View file

@ -1,19 +1,19 @@
use rustpython_parser::ast::{self, Expr, Stmt, StmtKind};
use rustpython_parser::ast::{self, Expr, Stmt};
pub fn name(stmt: &Stmt) -> &str {
match &stmt.node {
StmtKind::FunctionDef(ast::StmtFunctionDef { name, .. })
| StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef { name, .. }) => name.as_str(),
_ => panic!("Expected StmtKind::FunctionDef | StmtKind::AsyncFunctionDef"),
match stmt {
Stmt::FunctionDef(ast::StmtFunctionDef { name, .. })
| Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { name, .. }) => name.as_str(),
_ => panic!("Expected Stmt::FunctionDef | Stmt::AsyncFunctionDef"),
}
}
pub fn decorator_list(stmt: &Stmt) -> &Vec<Expr> {
match &stmt.node {
StmtKind::FunctionDef(ast::StmtFunctionDef { decorator_list, .. })
| StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef { decorator_list, .. }) => {
pub fn decorator_list(stmt: &Stmt) -> &[Expr] {
match stmt {
Stmt::FunctionDef(ast::StmtFunctionDef { decorator_list, .. })
| Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { decorator_list, .. }) => {
decorator_list
}
_ => panic!("Expected StmtKind::FunctionDef | StmtKind::AsyncFunctionDef"),
_ => panic!("Expected Stmt::FunctionDef | Stmt::AsyncFunctionDef"),
}
}

View file

@ -3,9 +3,8 @@
use num_bigint::BigInt;
use rustpython_parser::ast::{
self, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
ExcepthandlerKind, Expr, ExprContext, ExprKind, Identifier, Int, Keyword, MatchCase, Operator,
Pattern, PatternKind, Stmt, StmtKind, Unaryop, Withitem,
self, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler, Expr,
ExprContext, Identifier, Int, Keyword, MatchCase, Operator, Pattern, Stmt, Unaryop, Withitem,
};
#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
@ -136,8 +135,8 @@ pub struct ComparableAlias<'a> {
impl<'a> From<&'a Alias> for ComparableAlias<'a> {
fn from(alias: &'a Alias) -> Self {
Self {
name: alias.node.name.as_str(),
asname: alias.node.asname.as_deref(),
name: alias.name.as_str(),
asname: alias.asname.as_deref(),
}
}
}
@ -194,48 +193,50 @@ pub enum ComparablePattern<'a> {
impl<'a> From<&'a Pattern> for ComparablePattern<'a> {
fn from(pattern: &'a Pattern) -> Self {
match &pattern.node {
PatternKind::MatchValue(ast::PatternMatchValue { value }) => Self::MatchValue {
match pattern {
Pattern::MatchValue(ast::PatternMatchValue { value, .. }) => Self::MatchValue {
value: value.into(),
},
PatternKind::MatchSingleton(ast::PatternMatchSingleton { value }) => {
Pattern::MatchSingleton(ast::PatternMatchSingleton { value, .. }) => {
Self::MatchSingleton {
value: value.into(),
}
}
PatternKind::MatchSequence(ast::PatternMatchSequence { patterns }) => {
Pattern::MatchSequence(ast::PatternMatchSequence { patterns, .. }) => {
Self::MatchSequence {
patterns: patterns.iter().map(Into::into).collect(),
}
}
PatternKind::MatchMapping(ast::PatternMatchMapping {
Pattern::MatchMapping(ast::PatternMatchMapping {
keys,
patterns,
rest,
..
}) => Self::MatchMapping {
keys: keys.iter().map(Into::into).collect(),
patterns: patterns.iter().map(Into::into).collect(),
rest: rest.as_deref(),
},
PatternKind::MatchClass(ast::PatternMatchClass {
Pattern::MatchClass(ast::PatternMatchClass {
cls,
patterns,
kwd_attrs,
kwd_patterns,
..
}) => Self::MatchClass {
cls: cls.into(),
patterns: patterns.iter().map(Into::into).collect(),
kwd_attrs: kwd_attrs.iter().map(Identifier::as_str).collect(),
kwd_patterns: kwd_patterns.iter().map(Into::into).collect(),
},
PatternKind::MatchStar(ast::PatternMatchStar { name }) => Self::MatchStar {
Pattern::MatchStar(ast::PatternMatchStar { name, .. }) => Self::MatchStar {
name: name.as_deref(),
},
PatternKind::MatchAs(ast::PatternMatchAs { pattern, name }) => Self::MatchAs {
Pattern::MatchAs(ast::PatternMatchAs { pattern, name, .. }) => Self::MatchAs {
pattern: pattern.as_ref().map(Into::into),
name: name.as_deref(),
},
PatternKind::MatchOr(ast::PatternMatchOr { patterns }) => Self::MatchOr {
Pattern::MatchOr(ast::PatternMatchOr { patterns, .. }) => Self::MatchOr {
patterns: patterns.iter().map(Into::into).collect(),
},
}
@ -344,9 +345,9 @@ pub struct ComparableArg<'a> {
impl<'a> From<&'a Arg> for ComparableArg<'a> {
fn from(arg: &'a Arg) -> Self {
Self {
arg: arg.node.arg.as_str(),
annotation: arg.node.annotation.as_ref().map(Into::into),
type_comment: arg.node.type_comment.as_deref(),
arg: arg.arg.as_str(),
annotation: arg.annotation.as_ref().map(Into::into),
type_comment: arg.type_comment.as_deref(),
}
}
}
@ -360,8 +361,8 @@ pub struct ComparableKeyword<'a> {
impl<'a> From<&'a Keyword> for ComparableKeyword<'a> {
fn from(keyword: &'a Keyword) -> Self {
Self {
arg: keyword.node.arg.as_ref().map(Identifier::as_str),
value: (&keyword.node.value).into(),
arg: keyword.arg.as_ref().map(Identifier::as_str),
value: (&keyword.value).into(),
}
}
}
@ -396,8 +397,9 @@ pub enum ComparableExcepthandler<'a> {
impl<'a> From<&'a Excepthandler> for ComparableExcepthandler<'a> {
fn from(excepthandler: &'a Excepthandler) -> Self {
let ExcepthandlerKind::ExceptHandler(ast::ExcepthandlerExceptHandler { type_, name, body }) =
&excepthandler.node;
let Excepthandler::ExceptHandler(ast::ExcepthandlerExceptHandler {
type_, name, body, ..
}) = excepthandler;
Self::ExceptHandler {
type_: type_.as_ref().map(Into::into),
name: name.as_deref(),
@ -536,136 +538,230 @@ impl<'a> From<&'a Box<Expr>> for ComparableExpr<'a> {
impl<'a> From<&'a Expr> for ComparableExpr<'a> {
fn from(expr: &'a Expr) -> Self {
match &expr.node {
ExprKind::BoolOp(ast::ExprBoolOp { op, values }) => Self::BoolOp {
match expr {
Expr::BoolOp(ast::ExprBoolOp {
op,
values,
range: _range,
}) => Self::BoolOp {
op: op.into(),
values: values.iter().map(Into::into).collect(),
},
ExprKind::NamedExpr(ast::ExprNamedExpr { target, value }) => Self::NamedExpr {
Expr::NamedExpr(ast::ExprNamedExpr {
target,
value,
range: _range,
}) => Self::NamedExpr {
target: target.into(),
value: value.into(),
},
ExprKind::BinOp(ast::ExprBinOp { left, op, right }) => Self::BinOp {
Expr::BinOp(ast::ExprBinOp {
left,
op,
right,
range: _range,
}) => Self::BinOp {
left: left.into(),
op: op.into(),
right: right.into(),
},
ExprKind::UnaryOp(ast::ExprUnaryOp { op, operand }) => Self::UnaryOp {
Expr::UnaryOp(ast::ExprUnaryOp {
op,
operand,
range: _range,
}) => Self::UnaryOp {
op: op.into(),
operand: operand.into(),
},
ExprKind::Lambda(ast::ExprLambda { args, body }) => Self::Lambda {
Expr::Lambda(ast::ExprLambda {
args,
body,
range: _range,
}) => Self::Lambda {
args: (&**args).into(),
body: body.into(),
},
ExprKind::IfExp(ast::ExprIfExp { test, body, orelse }) => Self::IfExp {
Expr::IfExp(ast::ExprIfExp {
test,
body,
orelse,
range: _range,
}) => Self::IfExp {
test: test.into(),
body: body.into(),
orelse: orelse.into(),
},
ExprKind::Dict(ast::ExprDict { keys, values }) => Self::Dict {
Expr::Dict(ast::ExprDict {
keys,
values,
range: _range,
}) => Self::Dict {
keys: keys
.iter()
.map(|expr| expr.as_ref().map(Into::into))
.collect(),
values: values.iter().map(Into::into).collect(),
},
ExprKind::Set(ast::ExprSet { elts }) => Self::Set {
Expr::Set(ast::ExprSet {
elts,
range: _range,
}) => Self::Set {
elts: elts.iter().map(Into::into).collect(),
},
ExprKind::ListComp(ast::ExprListComp { elt, generators }) => Self::ListComp {
Expr::ListComp(ast::ExprListComp {
elt,
generators,
range: _range,
}) => Self::ListComp {
elt: elt.into(),
generators: generators.iter().map(Into::into).collect(),
},
ExprKind::SetComp(ast::ExprSetComp { elt, generators }) => Self::SetComp {
Expr::SetComp(ast::ExprSetComp {
elt,
generators,
range: _range,
}) => Self::SetComp {
elt: elt.into(),
generators: generators.iter().map(Into::into).collect(),
},
ExprKind::DictComp(ast::ExprDictComp {
Expr::DictComp(ast::ExprDictComp {
key,
value,
generators,
range: _range,
}) => Self::DictComp {
key: key.into(),
value: value.into(),
generators: generators.iter().map(Into::into).collect(),
},
ExprKind::GeneratorExp(ast::ExprGeneratorExp { elt, generators }) => {
Self::GeneratorExp {
elt: elt.into(),
generators: generators.iter().map(Into::into).collect(),
}
}
ExprKind::Await(ast::ExprAwait { value }) => Self::Await {
Expr::GeneratorExp(ast::ExprGeneratorExp {
elt,
generators,
range: _range,
}) => Self::GeneratorExp {
elt: elt.into(),
generators: generators.iter().map(Into::into).collect(),
},
Expr::Await(ast::ExprAwait {
value,
range: _range,
}) => Self::Await {
value: value.into(),
},
ExprKind::Yield(ast::ExprYield { value }) => Self::Yield {
Expr::Yield(ast::ExprYield {
value,
range: _range,
}) => Self::Yield {
value: value.as_ref().map(Into::into),
},
ExprKind::YieldFrom(ast::ExprYieldFrom { value }) => Self::YieldFrom {
Expr::YieldFrom(ast::ExprYieldFrom {
value,
range: _range,
}) => Self::YieldFrom {
value: value.into(),
},
ExprKind::Compare(ast::ExprCompare {
Expr::Compare(ast::ExprCompare {
left,
ops,
comparators,
range: _range,
}) => Self::Compare {
left: left.into(),
ops: ops.iter().map(Into::into).collect(),
comparators: comparators.iter().map(Into::into).collect(),
},
ExprKind::Call(ast::ExprCall {
Expr::Call(ast::ExprCall {
func,
args,
keywords,
range: _range,
}) => Self::Call {
func: func.into(),
args: args.iter().map(Into::into).collect(),
keywords: keywords.iter().map(Into::into).collect(),
},
ExprKind::FormattedValue(ast::ExprFormattedValue {
Expr::FormattedValue(ast::ExprFormattedValue {
value,
conversion,
format_spec,
range: _range,
}) => Self::FormattedValue {
value: value.into(),
conversion: *conversion,
format_spec: format_spec.as_ref().map(Into::into),
},
ExprKind::JoinedStr(ast::ExprJoinedStr { values }) => Self::JoinedStr {
Expr::JoinedStr(ast::ExprJoinedStr {
values,
range: _range,
}) => Self::JoinedStr {
values: values.iter().map(Into::into).collect(),
},
ExprKind::Constant(ast::ExprConstant { value, kind }) => Self::Constant {
Expr::Constant(ast::ExprConstant {
value,
kind,
range: _range,
}) => Self::Constant {
value: value.into(),
kind: kind.as_ref().map(String::as_str),
},
ExprKind::Attribute(ast::ExprAttribute { value, attr, ctx }) => Self::Attribute {
Expr::Attribute(ast::ExprAttribute {
value,
attr,
ctx,
range: _range,
}) => Self::Attribute {
value: value.into(),
attr: attr.as_str(),
ctx: ctx.into(),
},
ExprKind::Subscript(ast::ExprSubscript { value, slice, ctx }) => Self::Subscript {
Expr::Subscript(ast::ExprSubscript {
value,
slice,
ctx,
range: _range,
}) => Self::Subscript {
value: value.into(),
slice: slice.into(),
ctx: ctx.into(),
},
ExprKind::Starred(ast::ExprStarred { value, ctx }) => Self::Starred {
Expr::Starred(ast::ExprStarred {
value,
ctx,
range: _range,
}) => Self::Starred {
value: value.into(),
ctx: ctx.into(),
},
ExprKind::Name(ast::ExprName { id, ctx }) => Self::Name {
Expr::Name(ast::ExprName {
id,
ctx,
range: _range,
}) => Self::Name {
id: id.as_str(),
ctx: ctx.into(),
},
ExprKind::List(ast::ExprList { elts, ctx }) => Self::List {
Expr::List(ast::ExprList {
elts,
ctx,
range: _range,
}) => Self::List {
elts: elts.iter().map(Into::into).collect(),
ctx: ctx.into(),
},
ExprKind::Tuple(ast::ExprTuple { elts, ctx }) => Self::Tuple {
Expr::Tuple(ast::ExprTuple {
elts,
ctx,
range: _range,
}) => Self::Tuple {
elts: elts.iter().map(Into::into).collect(),
ctx: ctx.into(),
},
ExprKind::Slice(ast::ExprSlice { lower, upper, step }) => Self::Slice {
Expr::Slice(ast::ExprSlice {
lower,
upper,
step,
range: _range,
}) => Self::Slice {
lower: lower.as_ref().map(Into::into),
upper: upper.as_ref().map(Into::into),
step: step.as_ref().map(Into::into),
@ -803,14 +899,15 @@ pub enum ComparableStmt<'a> {
impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
fn from(stmt: &'a Stmt) -> Self {
match &stmt.node {
StmtKind::FunctionDef(ast::StmtFunctionDef {
match stmt {
Stmt::FunctionDef(ast::StmtFunctionDef {
name,
args,
body,
decorator_list,
returns,
type_comment,
range: _range,
}) => Self::FunctionDef {
name: name.as_str(),
args: args.into(),
@ -819,13 +916,14 @@ impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
returns: returns.as_ref().map(Into::into),
type_comment: type_comment.as_ref().map(std::string::String::as_str),
},
StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
name,
args,
body,
decorator_list,
returns,
type_comment,
range: _range,
}) => Self::AsyncFunctionDef {
name: name.as_str(),
args: args.into(),
@ -834,12 +932,13 @@ impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
returns: returns.as_ref().map(Into::into),
type_comment: type_comment.as_ref().map(std::string::String::as_str),
},
StmtKind::ClassDef(ast::StmtClassDef {
Stmt::ClassDef(ast::StmtClassDef {
name,
bases,
keywords,
body,
decorator_list,
range: _range,
}) => Self::ClassDef {
name: name.as_str(),
bases: bases.iter().map(Into::into).collect(),
@ -847,43 +946,57 @@ impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
body: body.iter().map(Into::into).collect(),
decorator_list: decorator_list.iter().map(Into::into).collect(),
},
StmtKind::Return(ast::StmtReturn { value }) => Self::Return {
Stmt::Return(ast::StmtReturn {
value,
range: _range,
}) => Self::Return {
value: value.as_ref().map(Into::into),
},
StmtKind::Delete(ast::StmtDelete { targets }) => Self::Delete {
Stmt::Delete(ast::StmtDelete {
targets,
range: _range,
}) => Self::Delete {
targets: targets.iter().map(Into::into).collect(),
},
StmtKind::Assign(ast::StmtAssign {
Stmt::Assign(ast::StmtAssign {
targets,
value,
type_comment,
range: _range,
}) => Self::Assign {
targets: targets.iter().map(Into::into).collect(),
value: value.into(),
type_comment: type_comment.as_ref().map(std::string::String::as_str),
},
StmtKind::AugAssign(ast::StmtAugAssign { target, op, value }) => Self::AugAssign {
Stmt::AugAssign(ast::StmtAugAssign {
target,
op,
value,
range: _range,
}) => Self::AugAssign {
target: target.into(),
op: op.into(),
value: value.into(),
},
StmtKind::AnnAssign(ast::StmtAnnAssign {
Stmt::AnnAssign(ast::StmtAnnAssign {
target,
annotation,
value,
simple,
range: _range,
}) => Self::AnnAssign {
target: target.into(),
annotation: annotation.into(),
value: value.as_ref().map(Into::into),
simple: *simple,
},
StmtKind::For(ast::StmtFor {
Stmt::For(ast::StmtFor {
target,
iter,
body,
orelse,
type_comment,
range: _range,
}) => Self::For {
target: target.into(),
iter: iter.into(),
@ -891,12 +1004,13 @@ impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
orelse: orelse.iter().map(Into::into).collect(),
type_comment: type_comment.as_ref().map(String::as_str),
},
StmtKind::AsyncFor(ast::StmtAsyncFor {
Stmt::AsyncFor(ast::StmtAsyncFor {
target,
iter,
body,
orelse,
type_comment,
range: _range,
}) => Self::AsyncFor {
target: target.into(),
iter: iter.into(),
@ -904,92 +1018,131 @@ impl<'a> From<&'a Stmt> for ComparableStmt<'a> {
orelse: orelse.iter().map(Into::into).collect(),
type_comment: type_comment.as_ref().map(String::as_str),
},
StmtKind::While(ast::StmtWhile { test, body, orelse }) => Self::While {
Stmt::While(ast::StmtWhile {
test,
body,
orelse,
range: _range,
}) => Self::While {
test: test.into(),
body: body.iter().map(Into::into).collect(),
orelse: orelse.iter().map(Into::into).collect(),
},
StmtKind::If(ast::StmtIf { test, body, orelse }) => Self::If {
Stmt::If(ast::StmtIf {
test,
body,
orelse,
range: _range,
}) => Self::If {
test: test.into(),
body: body.iter().map(Into::into).collect(),
orelse: orelse.iter().map(Into::into).collect(),
},
StmtKind::With(ast::StmtWith {
Stmt::With(ast::StmtWith {
items,
body,
type_comment,
range: _range,
}) => Self::With {
items: items.iter().map(Into::into).collect(),
body: body.iter().map(Into::into).collect(),
type_comment: type_comment.as_ref().map(String::as_str),
},
StmtKind::AsyncWith(ast::StmtAsyncWith {
Stmt::AsyncWith(ast::StmtAsyncWith {
items,
body,
type_comment,
range: _range,
}) => Self::AsyncWith {
items: items.iter().map(Into::into).collect(),
body: body.iter().map(Into::into).collect(),
type_comment: type_comment.as_ref().map(String::as_str),
},
StmtKind::Match(ast::StmtMatch { subject, cases }) => Self::Match {
Stmt::Match(ast::StmtMatch {
subject,
cases,
range: _range,
}) => Self::Match {
subject: subject.into(),
cases: cases.iter().map(Into::into).collect(),
},
StmtKind::Raise(ast::StmtRaise { exc, cause }) => Self::Raise {
Stmt::Raise(ast::StmtRaise {
exc,
cause,
range: _range,
}) => Self::Raise {
exc: exc.as_ref().map(Into::into),
cause: cause.as_ref().map(Into::into),
},
StmtKind::Try(ast::StmtTry {
Stmt::Try(ast::StmtTry {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => Self::Try {
body: body.iter().map(Into::into).collect(),
handlers: handlers.iter().map(Into::into).collect(),
orelse: orelse.iter().map(Into::into).collect(),
finalbody: finalbody.iter().map(Into::into).collect(),
},
StmtKind::TryStar(ast::StmtTryStar {
Stmt::TryStar(ast::StmtTryStar {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => Self::TryStar {
body: body.iter().map(Into::into).collect(),
handlers: handlers.iter().map(Into::into).collect(),
orelse: orelse.iter().map(Into::into).collect(),
finalbody: finalbody.iter().map(Into::into).collect(),
},
StmtKind::Assert(ast::StmtAssert { test, msg }) => Self::Assert {
Stmt::Assert(ast::StmtAssert {
test,
msg,
range: _range,
}) => Self::Assert {
test: test.into(),
msg: msg.as_ref().map(Into::into),
},
StmtKind::Import(ast::StmtImport { names }) => Self::Import {
Stmt::Import(ast::StmtImport {
names,
range: _range,
}) => Self::Import {
names: names.iter().map(Into::into).collect(),
},
StmtKind::ImportFrom(ast::StmtImportFrom {
Stmt::ImportFrom(ast::StmtImportFrom {
module,
names,
level,
range: _range,
}) => Self::ImportFrom {
module: module.as_deref(),
names: names.iter().map(Into::into).collect(),
level: *level,
},
StmtKind::Global(ast::StmtGlobal { names }) => Self::Global {
Stmt::Global(ast::StmtGlobal {
names,
range: _range,
}) => Self::Global {
names: names.iter().map(Identifier::as_str).collect(),
},
StmtKind::Nonlocal(ast::StmtNonlocal { names }) => Self::Nonlocal {
Stmt::Nonlocal(ast::StmtNonlocal {
names,
range: _range,
}) => Self::Nonlocal {
names: names.iter().map(Identifier::as_str).collect(),
},
StmtKind::Expr(ast::StmtExpr { value }) => Self::Expr {
Stmt::Expr(ast::StmtExpr {
value,
range: _range,
}) => Self::Expr {
value: value.into(),
},
StmtKind::Pass => Self::Pass,
StmtKind::Break => Self::Break,
StmtKind::Continue => Self::Continue,
Stmt::Pass(_) => Self::Pass,
Stmt::Break(_) => Self::Break,
Stmt::Continue(_) => Self::Continue,
}
}
}

File diff suppressed because it is too large Load diff

View file

@ -1,41 +1,61 @@
use ruff_text_size::TextRange;
use rustpython_parser::ast::{self, Expr, ExprKind, Keyword};
use rustpython_parser::ast::{self, Expr, Keyword};
fn relocate_keyword(keyword: &mut Keyword, location: TextRange) {
keyword.range = location;
relocate_expr(&mut keyword.node.value, location);
relocate_expr(&mut keyword.value, location);
}
/// Change an expression's location (recursively) to match a desired, fixed
/// location.
pub fn relocate_expr(expr: &mut Expr, location: TextRange) {
expr.range = location;
match &mut expr.node {
ExprKind::BoolOp(ast::ExprBoolOp { values, .. }) => {
match expr {
Expr::BoolOp(ast::ExprBoolOp { values, range, .. }) => {
*range = location;
for expr in values {
relocate_expr(expr, location);
}
}
ExprKind::NamedExpr(ast::ExprNamedExpr { target, value }) => {
Expr::NamedExpr(ast::ExprNamedExpr {
target,
value,
range,
}) => {
*range = location;
relocate_expr(target, location);
relocate_expr(value, location);
}
ExprKind::BinOp(ast::ExprBinOp { left, right, .. }) => {
Expr::BinOp(ast::ExprBinOp {
left, right, range, ..
}) => {
*range = location;
relocate_expr(left, location);
relocate_expr(right, location);
}
ExprKind::UnaryOp(ast::ExprUnaryOp { operand, .. }) => {
Expr::UnaryOp(ast::ExprUnaryOp { operand, range, .. }) => {
*range = location;
relocate_expr(operand, location);
}
ExprKind::Lambda(ast::ExprLambda { body, .. }) => {
Expr::Lambda(ast::ExprLambda { body, range, .. }) => {
*range = location;
relocate_expr(body, location);
}
ExprKind::IfExp(ast::ExprIfExp { test, body, orelse }) => {
Expr::IfExp(ast::ExprIfExp {
test,
body,
orelse,
range,
}) => {
*range = location;
relocate_expr(test, location);
relocate_expr(body, location);
relocate_expr(orelse, location);
}
ExprKind::Dict(ast::ExprDict { keys, values }) => {
Expr::Dict(ast::ExprDict {
keys,
values,
range,
}) => {
*range = location;
for expr in keys.iter_mut().flatten() {
relocate_expr(expr, location);
}
@ -43,44 +63,64 @@ pub fn relocate_expr(expr: &mut Expr, location: TextRange) {
relocate_expr(expr, location);
}
}
ExprKind::Set(ast::ExprSet { elts }) => {
Expr::Set(ast::ExprSet { elts, range }) => {
*range = location;
for expr in elts {
relocate_expr(expr, location);
}
}
ExprKind::ListComp(ast::ExprListComp { elt, .. }) => {
Expr::ListComp(ast::ExprListComp { elt, range, .. }) => {
*range = location;
relocate_expr(elt, location);
}
ExprKind::SetComp(ast::ExprSetComp { elt, .. }) => {
Expr::SetComp(ast::ExprSetComp { elt, range, .. }) => {
*range = location;
relocate_expr(elt, location);
}
ExprKind::DictComp(ast::ExprDictComp { key, value, .. }) => {
Expr::DictComp(ast::ExprDictComp {
key, value, range, ..
}) => {
*range = location;
relocate_expr(key, location);
relocate_expr(value, location);
}
ExprKind::GeneratorExp(ast::ExprGeneratorExp { elt, .. }) => {
Expr::GeneratorExp(ast::ExprGeneratorExp { elt, range, .. }) => {
*range = location;
relocate_expr(elt, location);
}
ExprKind::Await(ast::ExprAwait { value }) => relocate_expr(value, location),
ExprKind::Yield(ast::ExprYield { value }) => {
Expr::Await(ast::ExprAwait { value, range }) => {
*range = location;
relocate_expr(value, location);
}
Expr::Yield(ast::ExprYield { value, range }) => {
*range = location;
if let Some(expr) = value {
relocate_expr(expr, location);
}
}
ExprKind::YieldFrom(ast::ExprYieldFrom { value }) => relocate_expr(value, location),
ExprKind::Compare(ast::ExprCompare {
left, comparators, ..
Expr::YieldFrom(ast::ExprYieldFrom { value, range }) => {
*range = location;
relocate_expr(value, location);
}
Expr::Compare(ast::ExprCompare {
left,
comparators,
range,
..
}) => {
*range = location;
relocate_expr(left, location);
for expr in comparators {
relocate_expr(expr, location);
}
}
ExprKind::Call(ast::ExprCall {
Expr::Call(ast::ExprCall {
func,
args,
keywords,
range,
}) => {
*range = location;
relocate_expr(func, location);
for expr in args {
relocate_expr(expr, location);
@ -89,42 +129,67 @@ pub fn relocate_expr(expr: &mut Expr, location: TextRange) {
relocate_keyword(keyword, location);
}
}
ExprKind::FormattedValue(ast::ExprFormattedValue {
value, format_spec, ..
Expr::FormattedValue(ast::ExprFormattedValue {
value,
format_spec,
range,
..
}) => {
*range = location;
relocate_expr(value, location);
if let Some(expr) = format_spec {
relocate_expr(expr, location);
}
}
ExprKind::JoinedStr(ast::ExprJoinedStr { values }) => {
Expr::JoinedStr(ast::ExprJoinedStr { values, range }) => {
*range = location;
for expr in values {
relocate_expr(expr, location);
}
}
ExprKind::Constant(_) => {}
ExprKind::Attribute(ast::ExprAttribute { value, .. }) => {
Expr::Constant(ast::ExprConstant { range, .. }) => {
*range = location;
}
Expr::Attribute(ast::ExprAttribute { value, range, .. }) => {
*range = location;
relocate_expr(value, location);
}
ExprKind::Subscript(ast::ExprSubscript { value, slice, .. }) => {
Expr::Subscript(ast::ExprSubscript {
value,
slice,
range,
..
}) => {
*range = location;
relocate_expr(value, location);
relocate_expr(slice, location);
}
ExprKind::Starred(ast::ExprStarred { value, .. }) => {
Expr::Starred(ast::ExprStarred { value, range, .. }) => {
*range = location;
relocate_expr(value, location);
}
ExprKind::Name(_) => {}
ExprKind::List(ast::ExprList { elts, .. }) => {
Expr::Name(ast::ExprName { range, .. }) => {
*range = location;
}
Expr::List(ast::ExprList { elts, range, .. }) => {
*range = location;
for expr in elts {
relocate_expr(expr, location);
}
}
ExprKind::Tuple(ast::ExprTuple { elts, .. }) => {
Expr::Tuple(ast::ExprTuple { elts, range, .. }) => {
*range = location;
for expr in elts {
relocate_expr(expr, location);
}
}
ExprKind::Slice(ast::ExprSlice { lower, upper, step }) => {
Expr::Slice(ast::ExprSlice {
lower,
upper,
step,
range,
}) => {
*range = location;
if let Some(expr) = lower {
relocate_expr(expr, location);
}

View file

@ -5,8 +5,7 @@ use std::ops::Deref;
use rustpython_literal::escape::{AsciiEscape, Escape, UnicodeEscape};
use rustpython_parser::ast::{
self, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, ConversionFlag,
Excepthandler, ExcepthandlerKind, Expr, ExprKind, Identifier, Int, MatchCase, Operator,
Pattern, PatternKind, Stmt, StmtKind, Suite, Withitem,
Excepthandler, Expr, Identifier, Int, MatchCase, Operator, Pattern, Stmt, Suite, Withitem,
};
use crate::newlines::LineEnding;
@ -182,8 +181,8 @@ impl<'a> Generator<'a> {
}};
}
match &ast.node {
StmtKind::FunctionDef(ast::StmtFunctionDef {
match ast {
Stmt::FunctionDef(ast::StmtFunctionDef {
name,
args,
body,
@ -216,7 +215,7 @@ impl<'a> Generator<'a> {
self.newlines(2);
}
}
StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
name,
args,
body,
@ -248,12 +247,13 @@ impl<'a> Generator<'a> {
self.newlines(2);
}
}
StmtKind::ClassDef(ast::StmtClassDef {
Stmt::ClassDef(ast::StmtClassDef {
name,
bases,
keywords,
body,
decorator_list,
range: _range,
}) => {
self.newlines(if self.indent_depth == 0 { 2 } else { 1 });
statement!({
@ -274,13 +274,13 @@ impl<'a> Generator<'a> {
for keyword in keywords {
self.p_if(first, "(");
self.p_delim(&mut first, ", ");
if let Some(arg) = &keyword.node.arg {
if let Some(arg) = &keyword.arg {
self.p_id(arg);
self.p("=");
} else {
self.p("**");
}
self.unparse_expr(&keyword.node.value, precedence::MAX);
self.unparse_expr(&keyword.value, precedence::MAX);
}
self.p_if(!first, ")");
self.p(":");
@ -290,7 +290,10 @@ impl<'a> Generator<'a> {
self.newlines(2);
}
}
StmtKind::Return(ast::StmtReturn { value }) => {
Stmt::Return(ast::StmtReturn {
value,
range: _range,
}) => {
statement!({
if let Some(expr) = value {
self.p("return ");
@ -300,7 +303,10 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Delete(ast::StmtDelete { targets }) => {
Stmt::Delete(ast::StmtDelete {
targets,
range: _range,
}) => {
statement!({
self.p("del ");
let mut first = true;
@ -310,7 +316,7 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Assign(ast::StmtAssign { targets, value, .. }) => {
Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
statement!({
for target in targets {
self.unparse_expr(target, precedence::ASSIGN);
@ -319,7 +325,12 @@ impl<'a> Generator<'a> {
self.unparse_expr(value, precedence::ASSIGN);
});
}
StmtKind::AugAssign(ast::StmtAugAssign { target, op, value }) => {
Stmt::AugAssign(ast::StmtAugAssign {
target,
op,
value,
range: _range,
}) => {
statement!({
self.unparse_expr(target, precedence::AUG_ASSIGN);
self.p(" ");
@ -342,14 +353,15 @@ impl<'a> Generator<'a> {
self.unparse_expr(value, precedence::AUG_ASSIGN);
});
}
StmtKind::AnnAssign(ast::StmtAnnAssign {
Stmt::AnnAssign(ast::StmtAnnAssign {
target,
annotation,
value,
simple,
range: _range,
}) => {
statement!({
let need_parens = matches!(target.node, ExprKind::Name(_)) && !simple;
let need_parens = matches!(target.as_ref(), Expr::Name(_)) && !simple;
self.p_if(need_parens, "(");
self.unparse_expr(target, precedence::ANN_ASSIGN);
self.p_if(need_parens, ")");
@ -361,7 +373,7 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::For(ast::StmtFor {
Stmt::For(ast::StmtFor {
target,
iter,
body,
@ -383,7 +395,7 @@ impl<'a> Generator<'a> {
self.body(orelse);
}
}
StmtKind::AsyncFor(ast::StmtAsyncFor {
Stmt::AsyncFor(ast::StmtAsyncFor {
target,
iter,
body,
@ -405,7 +417,12 @@ impl<'a> Generator<'a> {
self.body(orelse);
}
}
StmtKind::While(ast::StmtWhile { test, body, orelse }) => {
Stmt::While(ast::StmtWhile {
test,
body,
orelse,
range: _range,
}) => {
statement!({
self.p("while ");
self.unparse_expr(test, precedence::WHILE);
@ -419,7 +436,12 @@ impl<'a> Generator<'a> {
self.body(orelse);
}
}
StmtKind::If(ast::StmtIf { test, body, orelse }) => {
Stmt::If(ast::StmtIf {
test,
body,
orelse,
range: _range,
}) => {
statement!({
self.p("if ");
self.unparse_expr(test, precedence::IF);
@ -427,10 +449,16 @@ impl<'a> Generator<'a> {
});
self.body(body);
let mut orelse_: &Vec<Stmt<U>> = orelse;
let mut orelse_: &[Stmt<U>] = orelse;
loop {
if orelse_.len() == 1 && matches!(orelse_[0].node, StmtKind::If(_)) {
if let StmtKind::If(ast::StmtIf { body, test, orelse }) = &orelse_[0].node {
if orelse_.len() == 1 && matches!(orelse_[0], Stmt::If(_)) {
if let Stmt::If(ast::StmtIf {
body,
test,
orelse,
range: _range,
}) = &orelse_[0]
{
statement!({
self.p("elif ");
self.unparse_expr(test, precedence::IF);
@ -450,7 +478,7 @@ impl<'a> Generator<'a> {
}
}
}
StmtKind::With(ast::StmtWith { items, body, .. }) => {
Stmt::With(ast::StmtWith { items, body, .. }) => {
statement!({
self.p("with ");
let mut first = true;
@ -462,7 +490,7 @@ impl<'a> Generator<'a> {
});
self.body(body);
}
StmtKind::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
Stmt::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
statement!({
self.p("async with ");
let mut first = true;
@ -474,7 +502,11 @@ impl<'a> Generator<'a> {
});
self.body(body);
}
StmtKind::Match(ast::StmtMatch { subject, cases }) => {
Stmt::Match(ast::StmtMatch {
subject,
cases,
range: _range,
}) => {
statement!({
self.p("match ");
self.unparse_expr(subject, precedence::MAX);
@ -488,7 +520,11 @@ impl<'a> Generator<'a> {
self.indent_depth -= 1;
}
}
StmtKind::Raise(ast::StmtRaise { exc, cause }) => {
Stmt::Raise(ast::StmtRaise {
exc,
cause,
range: _range,
}) => {
statement!({
self.p("raise");
if let Some(exc) = exc {
@ -501,11 +537,12 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Try(ast::StmtTry {
Stmt::Try(ast::StmtTry {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
statement!({
self.p("try:");
@ -531,11 +568,12 @@ impl<'a> Generator<'a> {
self.body(finalbody);
}
}
StmtKind::TryStar(ast::StmtTryStar {
Stmt::TryStar(ast::StmtTryStar {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
statement!({
self.p("try:");
@ -561,7 +599,11 @@ impl<'a> Generator<'a> {
self.body(finalbody);
}
}
StmtKind::Assert(ast::StmtAssert { test, msg }) => {
Stmt::Assert(ast::StmtAssert {
test,
msg,
range: _range,
}) => {
statement!({
self.p("assert ");
self.unparse_expr(test, precedence::ASSERT);
@ -571,7 +613,10 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Import(ast::StmtImport { names }) => {
Stmt::Import(ast::StmtImport {
names,
range: _range,
}) => {
statement!({
self.p("import ");
let mut first = true;
@ -581,10 +626,11 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::ImportFrom(ast::StmtImportFrom {
Stmt::ImportFrom(ast::StmtImportFrom {
module,
names,
level,
range: _range,
}) => {
statement!({
self.p("from ");
@ -602,7 +648,10 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Global(ast::StmtGlobal { names }) => {
Stmt::Global(ast::StmtGlobal {
names,
range: _range,
}) => {
statement!({
self.p("global ");
let mut first = true;
@ -612,7 +661,10 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Nonlocal(ast::StmtNonlocal { names }) => {
Stmt::Nonlocal(ast::StmtNonlocal {
names,
range: _range,
}) => {
statement!({
self.p("nonlocal ");
let mut first = true;
@ -622,22 +674,25 @@ impl<'a> Generator<'a> {
}
});
}
StmtKind::Expr(ast::StmtExpr { value }) => {
Stmt::Expr(ast::StmtExpr {
value,
range: _range,
}) => {
statement!({
self.unparse_expr(value, precedence::EXPR);
});
}
StmtKind::Pass => {
Stmt::Pass(_) => {
statement!({
self.p("pass");
});
}
StmtKind::Break => {
Stmt::Break(_) => {
statement!({
self.p("break");
});
}
StmtKind::Continue => {
Stmt::Continue(_) => {
statement!({
self.p("continue");
});
@ -646,11 +701,12 @@ impl<'a> Generator<'a> {
}
fn unparse_excepthandler<U>(&mut self, ast: &Excepthandler<U>, star: bool) {
match &ast.node {
ExcepthandlerKind::ExceptHandler(ast::ExcepthandlerExceptHandler {
match ast {
Excepthandler::ExceptHandler(ast::ExcepthandlerExceptHandler {
type_,
name,
body,
range: _range,
}) => {
self.p("except");
if star {
@ -671,14 +727,23 @@ impl<'a> Generator<'a> {
}
fn unparse_pattern<U>(&mut self, ast: &Pattern<U>) {
match &ast.node {
PatternKind::MatchValue(ast::PatternMatchValue { value }) => {
match ast {
Pattern::MatchValue(ast::PatternMatchValue {
value,
range: _range,
}) => {
self.unparse_expr(value, precedence::MAX);
}
PatternKind::MatchSingleton(ast::PatternMatchSingleton { value }) => {
Pattern::MatchSingleton(ast::PatternMatchSingleton {
value,
range: _range,
}) => {
self.unparse_constant(value);
}
PatternKind::MatchSequence(ast::PatternMatchSequence { patterns }) => {
Pattern::MatchSequence(ast::PatternMatchSequence {
patterns,
range: _range,
}) => {
self.p("[");
let mut first = true;
for pattern in patterns {
@ -687,10 +752,11 @@ impl<'a> Generator<'a> {
}
self.p("]");
}
PatternKind::MatchMapping(ast::PatternMatchMapping {
Pattern::MatchMapping(ast::PatternMatchMapping {
keys,
patterns,
rest,
range: _range,
}) => {
self.p("{");
let mut first = true;
@ -707,8 +773,11 @@ impl<'a> Generator<'a> {
}
self.p("}");
}
PatternKind::MatchClass(_) => {}
PatternKind::MatchStar(ast::PatternMatchStar { name }) => {
Pattern::MatchClass(_) => {}
Pattern::MatchStar(ast::PatternMatchStar {
name,
range: _range,
}) => {
self.p("*");
if let Some(name) = name {
self.p_id(name);
@ -716,7 +785,11 @@ impl<'a> Generator<'a> {
self.p("_");
}
}
PatternKind::MatchAs(ast::PatternMatchAs { pattern, name }) => {
Pattern::MatchAs(ast::PatternMatchAs {
pattern,
name,
range: _range,
}) => {
if let Some(pattern) = pattern {
self.unparse_pattern(pattern);
self.p(" as ");
@ -727,7 +800,10 @@ impl<'a> Generator<'a> {
self.p("_");
}
}
PatternKind::MatchOr(ast::PatternMatchOr { patterns }) => {
Pattern::MatchOr(ast::PatternMatchOr {
patterns,
range: _range,
}) => {
let mut first = true;
for pattern in patterns {
self.p_delim(&mut first, " | ");
@ -771,8 +847,12 @@ impl<'a> Generator<'a> {
ret
}};
}
match &ast.node {
ExprKind::BoolOp(ast::ExprBoolOp { op, values }) => {
match ast {
Expr::BoolOp(ast::ExprBoolOp {
op,
values,
range: _range,
}) => {
let (op, prec) = opprec!(bin, op, Boolop, And("and", AND), Or("or", OR));
group_if!(prec, {
let mut first = true;
@ -782,14 +862,23 @@ impl<'a> Generator<'a> {
}
});
}
ExprKind::NamedExpr(ast::ExprNamedExpr { target, value }) => {
Expr::NamedExpr(ast::ExprNamedExpr {
target,
value,
range: _range,
}) => {
group_if!(precedence::NAMED_EXPR, {
self.unparse_expr(target, precedence::NAMED_EXPR);
self.p(" := ");
self.unparse_expr(value, precedence::NAMED_EXPR + 1);
});
}
ExprKind::BinOp(ast::ExprBinOp { left, op, right }) => {
Expr::BinOp(ast::ExprBinOp {
left,
op,
right,
range: _range,
}) => {
let rassoc = matches!(op, Operator::Pow);
let (op, prec) = opprec!(
bin,
@ -815,7 +904,11 @@ impl<'a> Generator<'a> {
self.unparse_expr(right, prec + u8::from(!rassoc));
});
}
ExprKind::UnaryOp(ast::ExprUnaryOp { op, operand }) => {
Expr::UnaryOp(ast::ExprUnaryOp {
op,
operand,
range: _range,
}) => {
let (op, prec) = opprec!(
un,
op,
@ -830,7 +923,11 @@ impl<'a> Generator<'a> {
self.unparse_expr(operand, prec);
});
}
ExprKind::Lambda(ast::ExprLambda { args, body }) => {
Expr::Lambda(ast::ExprLambda {
args,
body,
range: _range,
}) => {
group_if!(precedence::LAMBDA, {
let npos = args.args.len() + args.posonlyargs.len();
self.p(if npos > 0 { "lambda " } else { "lambda" });
@ -839,7 +936,12 @@ impl<'a> Generator<'a> {
self.unparse_expr(body, precedence::LAMBDA);
});
}
ExprKind::IfExp(ast::ExprIfExp { test, body, orelse }) => {
Expr::IfExp(ast::ExprIfExp {
test,
body,
orelse,
range: _range,
}) => {
group_if!(precedence::IF_EXP, {
self.unparse_expr(body, precedence::IF_EXP + 1);
self.p(" if ");
@ -848,7 +950,11 @@ impl<'a> Generator<'a> {
self.unparse_expr(orelse, precedence::IF_EXP);
});
}
ExprKind::Dict(ast::ExprDict { keys, values }) => {
Expr::Dict(ast::ExprDict {
keys,
values,
range: _range,
}) => {
self.p("{");
let mut first = true;
for (k, v) in keys.iter().zip(values) {
@ -864,7 +970,10 @@ impl<'a> Generator<'a> {
}
self.p("}");
}
ExprKind::Set(ast::ExprSet { elts }) => {
Expr::Set(ast::ExprSet {
elts,
range: _range,
}) => {
if elts.is_empty() {
self.p("set()");
} else {
@ -877,22 +986,31 @@ impl<'a> Generator<'a> {
self.p("}");
}
}
ExprKind::ListComp(ast::ExprListComp { elt, generators }) => {
Expr::ListComp(ast::ExprListComp {
elt,
generators,
range: _range,
}) => {
self.p("[");
self.unparse_expr(elt, precedence::MAX);
self.unparse_comp(generators);
self.p("]");
}
ExprKind::SetComp(ast::ExprSetComp { elt, generators }) => {
Expr::SetComp(ast::ExprSetComp {
elt,
generators,
range: _range,
}) => {
self.p("{");
self.unparse_expr(elt, precedence::MAX);
self.unparse_comp(generators);
self.p("}");
}
ExprKind::DictComp(ast::ExprDictComp {
Expr::DictComp(ast::ExprDictComp {
key,
value,
generators,
range: _range,
}) => {
self.p("{");
self.unparse_expr(key, precedence::MAX);
@ -901,19 +1019,29 @@ impl<'a> Generator<'a> {
self.unparse_comp(generators);
self.p("}");
}
ExprKind::GeneratorExp(ast::ExprGeneratorExp { elt, generators }) => {
Expr::GeneratorExp(ast::ExprGeneratorExp {
elt,
generators,
range: _range,
}) => {
self.p("(");
self.unparse_expr(elt, precedence::COMMA);
self.unparse_comp(generators);
self.p(")");
}
ExprKind::Await(ast::ExprAwait { value }) => {
Expr::Await(ast::ExprAwait {
value,
range: _range,
}) => {
group_if!(precedence::AWAIT, {
self.p("await ");
self.unparse_expr(value, precedence::MAX);
});
}
ExprKind::Yield(ast::ExprYield { value }) => {
Expr::Yield(ast::ExprYield {
value,
range: _range,
}) => {
group_if!(precedence::YIELD, {
self.p("yield");
if let Some(value) = value {
@ -922,16 +1050,20 @@ impl<'a> Generator<'a> {
}
});
}
ExprKind::YieldFrom(ast::ExprYieldFrom { value }) => {
Expr::YieldFrom(ast::ExprYieldFrom {
value,
range: _range,
}) => {
group_if!(precedence::YIELD_FROM, {
self.p("yield from ");
self.unparse_expr(value, precedence::MAX);
});
}
ExprKind::Compare(ast::ExprCompare {
Expr::Compare(ast::ExprCompare {
left,
ops,
comparators,
range: _range,
}) => {
group_if!(precedence::CMP, {
let new_lvl = precedence::CMP + 1;
@ -954,18 +1086,20 @@ impl<'a> Generator<'a> {
}
});
}
ExprKind::Call(ast::ExprCall {
Expr::Call(ast::ExprCall {
func,
args,
keywords,
range: _range,
}) => {
self.unparse_expr(func, precedence::MAX);
self.p("(");
if let (
[Expr {
node: ExprKind::GeneratorExp(ast::ExprGeneratorExp { elt, generators }),
..
}],
[Expr::GeneratorExp(ast::ExprGeneratorExp {
elt,
generators,
range: _range,
})],
[],
) = (&**args, &**keywords)
{
@ -980,37 +1114,45 @@ impl<'a> Generator<'a> {
}
for kw in keywords {
self.p_delim(&mut first, ", ");
if let Some(arg) = &kw.node.arg {
if let Some(arg) = &kw.arg {
self.p_id(arg);
self.p("=");
self.unparse_expr(&kw.node.value, precedence::COMMA);
self.unparse_expr(&kw.value, precedence::COMMA);
} else {
self.p("**");
self.unparse_expr(&kw.node.value, precedence::MAX);
self.unparse_expr(&kw.value, precedence::MAX);
}
}
}
self.p(")");
}
ExprKind::FormattedValue(ast::ExprFormattedValue {
Expr::FormattedValue(ast::ExprFormattedValue {
value,
conversion,
format_spec,
range: _range,
}) => self.unparse_formatted(value, *conversion, format_spec.as_deref()),
ExprKind::JoinedStr(ast::ExprJoinedStr { values }) => {
Expr::JoinedStr(ast::ExprJoinedStr {
values,
range: _range,
}) => {
self.unparse_joinedstr(values, false);
}
ExprKind::Constant(ast::ExprConstant { value, kind }) => {
Expr::Constant(ast::ExprConstant {
value,
kind,
range: _range,
}) => {
if let Some(kind) = kind {
self.p(kind);
}
self.unparse_constant(value);
}
ExprKind::Attribute(ast::ExprAttribute { value, attr, .. }) => {
if let ExprKind::Constant(ast::ExprConstant {
Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
if let Expr::Constant(ast::ExprConstant {
value: Constant::Int(_),
..
}) = &value.node
}) = value.as_ref()
{
self.p("(");
self.unparse_expr(value, precedence::MAX);
@ -1021,18 +1163,18 @@ impl<'a> Generator<'a> {
};
self.p_id(attr);
}
ExprKind::Subscript(ast::ExprSubscript { value, slice, .. }) => {
Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
self.unparse_expr(value, precedence::MAX);
self.p("[");
self.unparse_expr(slice, precedence::SUBSCRIPT);
self.p("]");
}
ExprKind::Starred(ast::ExprStarred { value, .. }) => {
Expr::Starred(ast::ExprStarred { value, .. }) => {
self.p("*");
self.unparse_expr(value, precedence::MAX);
}
ExprKind::Name(ast::ExprName { id, .. }) => self.p_id(id),
ExprKind::List(ast::ExprList { elts, .. }) => {
Expr::Name(ast::ExprName { id, .. }) => self.p_id(id),
Expr::List(ast::ExprList { elts, .. }) => {
self.p("[");
let mut first = true;
for elt in elts {
@ -1041,7 +1183,7 @@ impl<'a> Generator<'a> {
}
self.p("]");
}
ExprKind::Tuple(ast::ExprTuple { elts, .. }) => {
Expr::Tuple(ast::ExprTuple { elts, .. }) => {
if elts.is_empty() {
self.p("()");
} else {
@ -1055,7 +1197,12 @@ impl<'a> Generator<'a> {
});
}
}
ExprKind::Slice(ast::ExprSlice { lower, upper, step }) => {
Expr::Slice(ast::ExprSlice {
lower,
upper,
step,
range: _range,
}) => {
if let Some(lower) = lower {
self.unparse_expr(lower, precedence::SLICE);
}
@ -1163,8 +1310,8 @@ impl<'a> Generator<'a> {
}
fn unparse_arg<U>(&mut self, arg: &Arg<U>) {
self.p_id(&arg.node.arg);
if let Some(ann) = &arg.node.annotation {
self.p_id(&arg.arg);
if let Some(ann) = &arg.annotation {
self.p(": ");
self.unparse_expr(ann, precedence::COMMA);
}
@ -1220,21 +1367,25 @@ impl<'a> Generator<'a> {
}
fn unparse_fstring_elem<U>(&mut self, expr: &Expr<U>, is_spec: bool) {
match &expr.node {
ExprKind::Constant(ast::ExprConstant { value, .. }) => {
match expr {
Expr::Constant(ast::ExprConstant { value, .. }) => {
if let Constant::Str(s) = value {
self.unparse_fstring_str(s);
} else {
unreachable!()
}
}
ExprKind::JoinedStr(ast::ExprJoinedStr { values }) => {
Expr::JoinedStr(ast::ExprJoinedStr {
values,
range: _range,
}) => {
self.unparse_joinedstr(values, is_spec);
}
ExprKind::FormattedValue(ast::ExprFormattedValue {
Expr::FormattedValue(ast::ExprFormattedValue {
value,
conversion,
format_spec,
range: _range,
}) => self.unparse_formatted(value, *conversion, format_spec.as_deref()),
_ => unreachable!(),
}
@ -1265,8 +1416,8 @@ impl<'a> Generator<'a> {
}
fn unparse_alias<U>(&mut self, alias: &Alias<U>) {
self.p_id(&alias.node.name);
if let Some(asname) = &alias.node.asname {
self.p_id(&alias.name);
if let Some(asname) = &alias.asname {
self.p(" as ");
self.p_id(asname);
}

View file

@ -1,6 +1,6 @@
//! Specialized AST visitor trait and walk functions that only visit statements.
use rustpython_parser::ast::{self, Excepthandler, ExcepthandlerKind, MatchCase, Stmt, StmtKind};
use rustpython_parser::ast::{self, Excepthandler, MatchCase, Stmt};
/// A trait for AST visitors that only need to visit statements.
pub trait StatementVisitor<'a> {
@ -25,48 +25,49 @@ pub fn walk_body<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, body: &'
}
pub fn walk_stmt<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
match &stmt.node {
StmtKind::FunctionDef(ast::StmtFunctionDef { body, .. }) => {
match stmt {
Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => {
visitor.visit_body(body);
}
StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. }) => {
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. }) => {
visitor.visit_body(body);
}
StmtKind::For(ast::StmtFor { body, orelse, .. }) => {
Stmt::For(ast::StmtFor { body, orelse, .. }) => {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::ClassDef(ast::StmtClassDef { body, .. }) => {
Stmt::ClassDef(ast::StmtClassDef { body, .. }) => {
visitor.visit_body(body);
}
StmtKind::AsyncFor(ast::StmtAsyncFor { body, orelse, .. }) => {
Stmt::AsyncFor(ast::StmtAsyncFor { body, orelse, .. }) => {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::While(ast::StmtWhile { body, orelse, .. }) => {
Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::If(ast::StmtIf { body, orelse, .. }) => {
Stmt::If(ast::StmtIf { body, orelse, .. }) => {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::With(ast::StmtWith { body, .. }) => {
Stmt::With(ast::StmtWith { body, .. }) => {
visitor.visit_body(body);
}
StmtKind::AsyncWith(ast::StmtAsyncWith { body, .. }) => {
Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => {
visitor.visit_body(body);
}
StmtKind::Match(ast::StmtMatch { cases, .. }) => {
Stmt::Match(ast::StmtMatch { cases, .. }) => {
for match_case in cases {
visitor.visit_match_case(match_case);
}
}
StmtKind::Try(ast::StmtTry {
Stmt::Try(ast::StmtTry {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
visitor.visit_body(body);
for excepthandler in handlers {
@ -75,11 +76,12 @@ pub fn walk_stmt<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, stmt: &'
visitor.visit_body(orelse);
visitor.visit_body(finalbody);
}
StmtKind::TryStar(ast::StmtTryStar {
Stmt::TryStar(ast::StmtTryStar {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
visitor.visit_body(body);
for excepthandler in handlers {
@ -96,8 +98,8 @@ pub fn walk_excepthandler<'a, V: StatementVisitor<'a> + ?Sized>(
visitor: &mut V,
excepthandler: &'a Excepthandler,
) {
match &excepthandler.node {
ExcepthandlerKind::ExceptHandler(ast::ExcepthandlerExceptHandler { body, .. }) => {
match excepthandler {
Excepthandler::ExceptHandler(ast::ExcepthandlerExceptHandler { body, .. }) => {
visitor.visit_body(body);
}
}

View file

@ -1,9 +1,8 @@
//! AST visitor trait and walk functions.
use rustpython_parser::ast::{
self, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
ExcepthandlerKind, Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern,
PatternKind, Stmt, StmtKind, Unaryop, Withitem,
self, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler, Expr,
ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, Unaryop, Withitem,
};
/// A trait for AST visitors. Visits all nodes in the AST recursively.
@ -80,8 +79,8 @@ pub fn walk_body<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, body: &'a [Stmt])
}
pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
match &stmt.node {
StmtKind::FunctionDef(ast::StmtFunctionDef {
match stmt {
Stmt::FunctionDef(ast::StmtFunctionDef {
args,
body,
decorator_list,
@ -97,7 +96,7 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
}
visitor.visit_body(body);
}
StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
args,
body,
decorator_list,
@ -113,7 +112,7 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
}
visitor.visit_body(body);
}
StmtKind::ClassDef(ast::StmtClassDef {
Stmt::ClassDef(ast::StmtClassDef {
bases,
keywords,
body,
@ -131,28 +130,39 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
}
visitor.visit_body(body);
}
StmtKind::Return(ast::StmtReturn { value }) => {
Stmt::Return(ast::StmtReturn {
value,
range: _range,
}) => {
if let Some(expr) = value {
visitor.visit_expr(expr);
}
}
StmtKind::Delete(ast::StmtDelete { targets }) => {
Stmt::Delete(ast::StmtDelete {
targets,
range: _range,
}) => {
for expr in targets {
visitor.visit_expr(expr);
}
}
StmtKind::Assign(ast::StmtAssign { targets, value, .. }) => {
Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
visitor.visit_expr(value);
for expr in targets {
visitor.visit_expr(expr);
}
}
StmtKind::AugAssign(ast::StmtAugAssign { target, op, value }) => {
Stmt::AugAssign(ast::StmtAugAssign {
target,
op,
value,
range: _range,
}) => {
visitor.visit_expr(target);
visitor.visit_operator(op);
visitor.visit_expr(value);
}
StmtKind::AnnAssign(ast::StmtAnnAssign {
Stmt::AnnAssign(ast::StmtAnnAssign {
target,
annotation,
value,
@ -164,7 +174,7 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
}
visitor.visit_expr(target);
}
StmtKind::For(ast::StmtFor {
Stmt::For(ast::StmtFor {
target,
iter,
body,
@ -176,7 +186,7 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::AsyncFor(ast::StmtAsyncFor {
Stmt::AsyncFor(ast::StmtAsyncFor {
target,
iter,
body,
@ -188,35 +198,53 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::While(ast::StmtWhile { test, body, orelse }) => {
Stmt::While(ast::StmtWhile {
test,
body,
orelse,
range: _range,
}) => {
visitor.visit_expr(test);
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::If(ast::StmtIf { test, body, orelse }) => {
Stmt::If(ast::StmtIf {
test,
body,
orelse,
range: _range,
}) => {
visitor.visit_expr(test);
visitor.visit_body(body);
visitor.visit_body(orelse);
}
StmtKind::With(ast::StmtWith { items, body, .. }) => {
Stmt::With(ast::StmtWith { items, body, .. }) => {
for withitem in items {
visitor.visit_withitem(withitem);
}
visitor.visit_body(body);
}
StmtKind::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
Stmt::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
for withitem in items {
visitor.visit_withitem(withitem);
}
visitor.visit_body(body);
}
StmtKind::Match(ast::StmtMatch { subject, cases }) => {
Stmt::Match(ast::StmtMatch {
subject,
cases,
range: _range,
}) => {
visitor.visit_expr(subject);
for match_case in cases {
visitor.visit_match_case(match_case);
}
}
StmtKind::Raise(ast::StmtRaise { exc, cause }) => {
Stmt::Raise(ast::StmtRaise {
exc,
cause,
range: _range,
}) => {
if let Some(expr) = exc {
visitor.visit_expr(expr);
};
@ -224,11 +252,12 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
visitor.visit_expr(expr);
};
}
StmtKind::Try(ast::StmtTry {
Stmt::Try(ast::StmtTry {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
visitor.visit_body(body);
for excepthandler in handlers {
@ -237,11 +266,12 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
visitor.visit_body(orelse);
visitor.visit_body(finalbody);
}
StmtKind::TryStar(ast::StmtTryStar {
Stmt::TryStar(ast::StmtTryStar {
body,
handlers,
orelse,
finalbody,
range: _range,
}) => {
visitor.visit_body(body);
for excepthandler in handlers {
@ -250,62 +280,100 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
visitor.visit_body(orelse);
visitor.visit_body(finalbody);
}
StmtKind::Assert(ast::StmtAssert { test, msg }) => {
Stmt::Assert(ast::StmtAssert {
test,
msg,
range: _range,
}) => {
visitor.visit_expr(test);
if let Some(expr) = msg {
visitor.visit_expr(expr);
}
}
StmtKind::Import(ast::StmtImport { names }) => {
Stmt::Import(ast::StmtImport {
names,
range: _range,
}) => {
for alias in names {
visitor.visit_alias(alias);
}
}
StmtKind::ImportFrom(ast::StmtImportFrom { names, .. }) => {
Stmt::ImportFrom(ast::StmtImportFrom { names, .. }) => {
for alias in names {
visitor.visit_alias(alias);
}
}
StmtKind::Global(_) => {}
StmtKind::Nonlocal(_) => {}
StmtKind::Expr(ast::StmtExpr { value }) => visitor.visit_expr(value),
StmtKind::Pass => {}
StmtKind::Break => {}
StmtKind::Continue => {}
Stmt::Global(_) => {}
Stmt::Nonlocal(_) => {}
Stmt::Expr(ast::StmtExpr {
value,
range: _range,
}) => visitor.visit_expr(value),
Stmt::Pass(_) | Stmt::Break(_) | Stmt::Continue(_) => {}
}
}
pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
match &expr.node {
ExprKind::BoolOp(ast::ExprBoolOp { op, values }) => {
match expr {
Expr::BoolOp(ast::ExprBoolOp {
op,
values,
range: _range,
}) => {
visitor.visit_boolop(op);
for expr in values {
visitor.visit_expr(expr);
}
}
ExprKind::NamedExpr(ast::ExprNamedExpr { target, value }) => {
Expr::NamedExpr(ast::ExprNamedExpr {
target,
value,
range: _range,
}) => {
visitor.visit_expr(value);
visitor.visit_expr(target);
}
ExprKind::BinOp(ast::ExprBinOp { left, op, right }) => {
Expr::BinOp(ast::ExprBinOp {
left,
op,
right,
range: _range,
}) => {
visitor.visit_expr(left);
visitor.visit_operator(op);
visitor.visit_expr(right);
}
ExprKind::UnaryOp(ast::ExprUnaryOp { op, operand }) => {
Expr::UnaryOp(ast::ExprUnaryOp {
op,
operand,
range: _range,
}) => {
visitor.visit_unaryop(op);
visitor.visit_expr(operand);
}
ExprKind::Lambda(ast::ExprLambda { args, body }) => {
Expr::Lambda(ast::ExprLambda {
args,
body,
range: _range,
}) => {
visitor.visit_arguments(args);
visitor.visit_expr(body);
}
ExprKind::IfExp(ast::ExprIfExp { test, body, orelse }) => {
Expr::IfExp(ast::ExprIfExp {
test,
body,
orelse,
range: _range,
}) => {
visitor.visit_expr(test);
visitor.visit_expr(body);
visitor.visit_expr(orelse);
}
ExprKind::Dict(ast::ExprDict { keys, values }) => {
Expr::Dict(ast::ExprDict {
keys,
values,
range: _range,
}) => {
for expr in keys.iter().flatten() {
visitor.visit_expr(expr);
}
@ -313,27 +381,39 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(expr);
}
}
ExprKind::Set(ast::ExprSet { elts }) => {
Expr::Set(ast::ExprSet {
elts,
range: _range,
}) => {
for expr in elts {
visitor.visit_expr(expr);
}
}
ExprKind::ListComp(ast::ExprListComp { elt, generators }) => {
Expr::ListComp(ast::ExprListComp {
elt,
generators,
range: _range,
}) => {
for comprehension in generators {
visitor.visit_comprehension(comprehension);
}
visitor.visit_expr(elt);
}
ExprKind::SetComp(ast::ExprSetComp { elt, generators }) => {
Expr::SetComp(ast::ExprSetComp {
elt,
generators,
range: _range,
}) => {
for comprehension in generators {
visitor.visit_comprehension(comprehension);
}
visitor.visit_expr(elt);
}
ExprKind::DictComp(ast::ExprDictComp {
Expr::DictComp(ast::ExprDictComp {
key,
value,
generators,
range: _range,
}) => {
for comprehension in generators {
visitor.visit_comprehension(comprehension);
@ -341,23 +421,37 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(key);
visitor.visit_expr(value);
}
ExprKind::GeneratorExp(ast::ExprGeneratorExp { elt, generators }) => {
Expr::GeneratorExp(ast::ExprGeneratorExp {
elt,
generators,
range: _range,
}) => {
for comprehension in generators {
visitor.visit_comprehension(comprehension);
}
visitor.visit_expr(elt);
}
ExprKind::Await(ast::ExprAwait { value }) => visitor.visit_expr(value),
ExprKind::Yield(ast::ExprYield { value }) => {
Expr::Await(ast::ExprAwait {
value,
range: _range,
}) => visitor.visit_expr(value),
Expr::Yield(ast::ExprYield {
value,
range: _range,
}) => {
if let Some(expr) = value {
visitor.visit_expr(expr);
}
}
ExprKind::YieldFrom(ast::ExprYieldFrom { value }) => visitor.visit_expr(value),
ExprKind::Compare(ast::ExprCompare {
Expr::YieldFrom(ast::ExprYieldFrom {
value,
range: _range,
}) => visitor.visit_expr(value),
Expr::Compare(ast::ExprCompare {
left,
ops,
comparators,
range: _range,
}) => {
visitor.visit_expr(left);
for cmpop in ops {
@ -367,10 +461,11 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(expr);
}
}
ExprKind::Call(ast::ExprCall {
Expr::Call(ast::ExprCall {
func,
args,
keywords,
range: _range,
}) => {
visitor.visit_expr(func);
for expr in args {
@ -380,7 +475,7 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_keyword(keyword);
}
}
ExprKind::FormattedValue(ast::ExprFormattedValue {
Expr::FormattedValue(ast::ExprFormattedValue {
value, format_spec, ..
}) => {
visitor.visit_expr(value);
@ -388,41 +483,66 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_format_spec(expr);
}
}
ExprKind::JoinedStr(ast::ExprJoinedStr { values }) => {
Expr::JoinedStr(ast::ExprJoinedStr {
values,
range: _range,
}) => {
for expr in values {
visitor.visit_expr(expr);
}
}
ExprKind::Constant(ast::ExprConstant { value, .. }) => visitor.visit_constant(value),
ExprKind::Attribute(ast::ExprAttribute { value, ctx, .. }) => {
Expr::Constant(ast::ExprConstant { value, .. }) => visitor.visit_constant(value),
Expr::Attribute(ast::ExprAttribute { value, ctx, .. }) => {
visitor.visit_expr(value);
visitor.visit_expr_context(ctx);
}
ExprKind::Subscript(ast::ExprSubscript { value, slice, ctx }) => {
Expr::Subscript(ast::ExprSubscript {
value,
slice,
ctx,
range: _range,
}) => {
visitor.visit_expr(value);
visitor.visit_expr(slice);
visitor.visit_expr_context(ctx);
}
ExprKind::Starred(ast::ExprStarred { value, ctx }) => {
Expr::Starred(ast::ExprStarred {
value,
ctx,
range: _range,
}) => {
visitor.visit_expr(value);
visitor.visit_expr_context(ctx);
}
ExprKind::Name(ast::ExprName { ctx, .. }) => {
Expr::Name(ast::ExprName { ctx, .. }) => {
visitor.visit_expr_context(ctx);
}
ExprKind::List(ast::ExprList { elts, ctx }) => {
Expr::List(ast::ExprList {
elts,
ctx,
range: _range,
}) => {
for expr in elts {
visitor.visit_expr(expr);
}
visitor.visit_expr_context(ctx);
}
ExprKind::Tuple(ast::ExprTuple { elts, ctx }) => {
Expr::Tuple(ast::ExprTuple {
elts,
ctx,
range: _range,
}) => {
for expr in elts {
visitor.visit_expr(expr);
}
visitor.visit_expr_context(ctx);
}
ExprKind::Slice(ast::ExprSlice { lower, upper, step }) => {
Expr::Slice(ast::ExprSlice {
lower,
upper,
step,
range: _range,
}) => {
if let Some(expr) = lower {
visitor.visit_expr(expr);
}
@ -459,10 +579,8 @@ pub fn walk_excepthandler<'a, V: Visitor<'a> + ?Sized>(
visitor: &mut V,
excepthandler: &'a Excepthandler,
) {
match &excepthandler.node {
ExcepthandlerKind::ExceptHandler(ast::ExcepthandlerExceptHandler {
type_, body, ..
}) => {
match excepthandler {
Excepthandler::ExceptHandler(ast::ExcepthandlerExceptHandler { type_, body, .. }) => {
if let Some(expr) = type_ {
visitor.visit_expr(expr);
}
@ -496,13 +614,13 @@ pub fn walk_arguments<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, arguments: &
}
pub fn walk_arg<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, arg: &'a Arg) {
if let Some(expr) = &arg.node.annotation {
if let Some(expr) = &arg.annotation {
visitor.visit_annotation(expr);
}
}
pub fn walk_keyword<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, keyword: &'a Keyword) {
visitor.visit_expr(&keyword.node.value);
visitor.visit_expr(&keyword.value);
}
pub fn walk_withitem<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, withitem: &'a Withitem) {
@ -521,17 +639,26 @@ pub fn walk_match_case<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, match_case:
}
pub fn walk_pattern<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, pattern: &'a Pattern) {
match &pattern.node {
PatternKind::MatchValue(ast::PatternMatchValue { value }) => visitor.visit_expr(value),
PatternKind::MatchSingleton(ast::PatternMatchSingleton { value }) => {
match pattern {
Pattern::MatchValue(ast::PatternMatchValue {
value,
range: _range,
}) => visitor.visit_expr(value),
Pattern::MatchSingleton(ast::PatternMatchSingleton {
value,
range: _range,
}) => {
visitor.visit_constant(value);
}
PatternKind::MatchSequence(ast::PatternMatchSequence { patterns }) => {
Pattern::MatchSequence(ast::PatternMatchSequence {
patterns,
range: _range,
}) => {
for pattern in patterns {
visitor.visit_pattern(pattern);
}
}
PatternKind::MatchMapping(ast::PatternMatchMapping { keys, patterns, .. }) => {
Pattern::MatchMapping(ast::PatternMatchMapping { keys, patterns, .. }) => {
for expr in keys {
visitor.visit_expr(expr);
}
@ -539,7 +666,7 @@ pub fn walk_pattern<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, pattern: &'a P
visitor.visit_pattern(pattern);
}
}
PatternKind::MatchClass(ast::PatternMatchClass {
Pattern::MatchClass(ast::PatternMatchClass {
cls,
patterns,
kwd_patterns,
@ -554,13 +681,16 @@ pub fn walk_pattern<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, pattern: &'a P
visitor.visit_pattern(pattern);
}
}
PatternKind::MatchStar(_) => {}
PatternKind::MatchAs(ast::PatternMatchAs { pattern, .. }) => {
Pattern::MatchStar(_) => {}
Pattern::MatchAs(ast::PatternMatchAs { pattern, .. }) => {
if let Some(pattern) = pattern {
visitor.visit_pattern(pattern);
}
}
PatternKind::MatchOr(ast::PatternMatchOr { patterns }) => {
Pattern::MatchOr(ast::PatternMatchOr {
patterns,
range: _range,
}) => {
for pattern in patterns {
visitor.visit_pattern(pattern);
}

View file

@ -1,10 +1,13 @@
use ruff_text_size::TextRange;
use rustpython_parser::ast::Attributed;
use rustpython_parser::ast::Ranged;
use crate::source_code::Locator;
/// Extract the leading indentation from a line.
pub fn indentation<'a, T>(locator: &'a Locator, located: &Attributed<T>) -> Option<&'a str> {
pub fn indentation<'a, T>(locator: &'a Locator, located: &T) -> Option<&'a str>
where
T: Ranged,
{
let line_start = locator.line_start(located.start());
let indentation = &locator.contents()[TextRange::new(line_start, located.start())];