Skip to content

Commit 569bee1

Browse files
ShaharNavehShaharNaveh
authored
Use ruff_python_ast::visitor::Visitor for detecting await (#6466)
1 parent 8b3bf55 commit 569bee1

1 file changed

Lines changed: 30 additions & 128 deletions

File tree

‎crates/codegen/src/compile.rs‎

Lines changed: 30 additions & 128 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,13 @@ use ruff_python_ast::{
2424
Alias, Arguments, BoolOp, CmpOp, Comprehension, ConversionFlag, DebugText, Decorator, DictItem,
2525
ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprContext,
2626
ExprFString, ExprList, ExprName, ExprSlice, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp,
27-
FString, FStringFlags, FStringPart, Identifier, Int, InterpolatedElement,
28-
InterpolatedStringElement, InterpolatedStringElements, Keyword, MatchCase, ModExpression,
29-
ModModule, Operator, Parameters, Pattern, PatternMatchAs, PatternMatchClass,
30-
PatternMatchMapping, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton,
31-
PatternMatchStar, PatternMatchValue, Singleton, Stmt, StmtExpr, TypeParam, TypeParamParamSpec,
32-
TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams, UnaryOp, WithItem,
27+
FString, FStringFlags, FStringPart, Identifier, Int, InterpolatedStringElement,
28+
InterpolatedStringElements, Keyword, MatchCase, ModExpression, ModModule, Operator, Parameters,
29+
Pattern, PatternMatchAs, PatternMatchClass, PatternMatchMapping, PatternMatchOr,
30+
PatternMatchSequence, PatternMatchSingleton, PatternMatchStar, PatternMatchValue, Singleton,
31+
Stmt, StmtExpr, TypeParam, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple,
32+
TypeParams, UnaryOp, WithItem,
33+
visitor::{Visitor, walk_expr},
3334
};
3435
use ruff_text_size::{Ranged, TextRange};
3536
use rustpython_compiler_core::{
@@ -5435,134 +5436,35 @@ impl Compiler {
54355436

54365437
/// Whether the expression contains an await expression and
54375438
/// thus requires the function to be async.
5438-
/// Async with and async for are statements, so I won't check for them here
5439+
///
5440+
/// Both:
5441+
/// ```py
5442+
/// async with: ...
5443+
/// async for: ...
5444+
/// ```
5445+
/// are statements, so we won't check for them here
54395446
fn contains_await(expression: &Expr) -> bool {
5440-
use ruff_python_ast::*;
5447+
#[derive(Default)]
5448+
struct AwaitVisitor {
5449+
found: bool,
5450+
}
54415451

5442-
match &expression {
5443-
Expr::Call(ExprCall {
5444-
func, arguments, ..
5445-
}) => {
5446-
Self::contains_await(func)
5447-
|| arguments.args.iter().any(Self::contains_await)
5448-
|| arguments
5449-
.keywords
5450-
.iter()
5451-
.any(|kw| Self::contains_await(&kw.value))
5452-
}
5453-
Expr::BoolOp(ExprBoolOp { values, .. }) => values.iter().any(Self::contains_await),
5454-
Expr::BinOp(ExprBinOp { left, right, .. }) => {
5455-
Self::contains_await(left) || Self::contains_await(right)
5456-
}
5457-
Expr::Subscript(ExprSubscript { value, slice, .. }) => {
5458-
Self::contains_await(value) || Self::contains_await(slice)
5459-
}
5460-
Expr::UnaryOp(ExprUnaryOp { operand, .. }) => Self::contains_await(operand),
5461-
Expr::Attribute(ExprAttribute { value, .. }) => Self::contains_await(value),
5462-
Expr::Compare(ExprCompare {
5463-
left, comparators, ..
5464-
}) => Self::contains_await(left) || comparators.iter().any(Self::contains_await),
5465-
Expr::List(ExprList { elts, .. }) => elts.iter().any(Self::contains_await),
5466-
Expr::Tuple(ExprTuple { elts, .. }) => elts.iter().any(Self::contains_await),
5467-
Expr::Set(ExprSet { elts, .. }) => elts.iter().any(Self::contains_await),
5468-
Expr::Dict(ExprDict { items, .. }) => items
5469-
.iter()
5470-
.flat_map(|item| &item.key)
5471-
.any(Self::contains_await),
5472-
Expr::Slice(ExprSlice {
5473-
lower, upper, step, ..
5474-
}) => {
5475-
lower.as_deref().is_some_and(Self::contains_await)
5476-
|| upper.as_deref().is_some_and(Self::contains_await)
5477-
|| step.as_deref().is_some_and(Self::contains_await)
5478-
}
5479-
Expr::Yield(ExprYield { value, .. }) => {
5480-
value.as_deref().is_some_and(Self::contains_await)
5481-
}
5482-
Expr::Await(ExprAwait { .. }) => true,
5483-
Expr::YieldFrom(ExprYieldFrom { value, .. }) => Self::contains_await(value),
5484-
Expr::Name(ExprName { .. }) => false,
5485-
Expr::Lambda(ExprLambda { body, .. }) => Self::contains_await(body),
5486-
Expr::ListComp(ExprListComp {
5487-
elt, generators, ..
5488-
}) => {
5489-
Self::contains_await(elt)
5490-
|| generators.iter().any(|jen| Self::contains_await(&jen.iter))
5491-
}
5492-
Expr::SetComp(ExprSetComp {
5493-
elt, generators, ..
5494-
}) => {
5495-
Self::contains_await(elt)
5496-
|| generators.iter().any(|jen| Self::contains_await(&jen.iter))
5497-
}
5498-
Expr::DictComp(ExprDictComp {
5499-
key,
5500-
value,
5501-
generators,
5502-
..
5503-
}) => {
5504-
Self::contains_await(key)
5505-
|| Self::contains_await(value)
5506-
|| generators.iter().any(|jen| Self::contains_await(&jen.iter))
5507-
}
5508-
Expr::Generator(ExprGenerator {
5509-
elt, generators, ..
5510-
}) => {
5511-
Self::contains_await(elt)
5512-
|| generators.iter().any(|jen| Self::contains_await(&jen.iter))
5513-
}
5514-
Expr::Starred(expr) => Self::contains_await(&expr.value),
5515-
Expr::If(ExprIf {
5516-
test, body, orelse, ..
5517-
}) => {
5518-
Self::contains_await(test)
5519-
|| Self::contains_await(body)
5520-
|| Self::contains_await(orelse)
5521-
}
5452+
impl Visitor<'_> for AwaitVisitor {
5453+
fn visit_expr(&mut self, expr: &Expr) {
5454+
if self.found {
5455+
return;
5456+
}
55225457

5523-
Expr::Named(ExprNamed {
5524-
target,
5525-
value,
5526-
node_index: _,
5527-
range: _,
5528-
}) => Self::contains_await(target) || Self::contains_await(value),
5529-
Expr::FString(fstring) => {
5530-
Self::interpolated_string_contains_await(fstring.value.elements())
5531-
}
5532-
Expr::TString(tstring) => {
5533-
Self::interpolated_string_contains_await(tstring.value.elements())
5458+
match expr {
5459+
Expr::Await(_) => self.found = true,
5460+
_ => walk_expr(self, expr),
5461+
}
55345462
}
5535-
Expr::StringLiteral(_)
5536-
| Expr::BytesLiteral(_)
5537-
| Expr::NumberLiteral(_)
5538-
| Expr::BooleanLiteral(_)
5539-
| Expr::NoneLiteral(_)
5540-
| Expr::EllipsisLiteral(_)
5541-
| Expr::IpyEscapeCommand(_) => false,
5542-
}
5543-
}
5544-
5545-
fn interpolated_string_contains_await<'a>(
5546-
mut elements: impl Iterator<Item = &'a InterpolatedStringElement>,
5547-
) -> bool {
5548-
fn interpolated_element_contains_await<F: Copy + Fn(&Expr) -> bool>(
5549-
expr_element: &InterpolatedElement,
5550-
contains_await: F,
5551-
) -> bool {
5552-
contains_await(&expr_element.expression)
5553-
|| expr_element
5554-
.format_spec
5555-
.iter()
5556-
.flat_map(|spec| spec.elements.interpolations())
5557-
.any(|element| interpolated_element_contains_await(element, contains_await))
55585463
}
55595464

5560-
elements.any(|element| match element {
5561-
InterpolatedStringElement::Interpolation(expr_element) => {
5562-
interpolated_element_contains_await(expr_element, Self::contains_await)
5563-
}
5564-
InterpolatedStringElement::Literal(_) => false,
5565-
})
5465+
let mut visitor = AwaitVisitor::default();
5466+
visitor.visit_expr(expression);
5467+
visitor.found
55665468
}
55675469

55685470
fn compile_expr_fstring(&mut self, fstring: &ExprFString) -> CompileResult<()> {

0 commit comments

Comments
 (0)