mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-07-23 04:55:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

Remove async AST node variants for with, for, and def (#6369)

Browse source 
## Summary

Per the suggestion in
https://github.com/astral-sh/ruff/discussions/6183, this PR removes
`AsyncWith`, `AsyncFor`, and `AsyncFunctionDef`, replacing them with an
`is_async` field on the non-async variants of those structs. Unlike an
interpreter, we _generally_ have identical handling for these nodes, so
separating them into distinct variants adds complexity from which we
don't really benefit. This can be seen below, where we get to remove a
_ton_ of code related to adding generic `Any*` wrappers, and a ton of
duplicate branches for these cases.

## Test Plan

`cargo test` is unchanged, apart from parser snapshots.
This commit is contained in:
Charlie Marsh 2023-08-07 12:36:02 -04:00 • committed by GitHub
parent c895252aae
commit daefa74e9a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
91 changed files with 375 additions and 1478 deletions
Download patch file Download diff file Expand all files Collapse all files

14
crates/ruff_python_formatter/src/comments/placement.rs
View file

@ -67,9 +67,7 @@ pub(super) fn place_comment<'a>(
handle_module_level_own_line_comment_before_class_or_function_comment(comment, locator)
}
AnyNodeRef::WithItem(_) => handle_with_item_comment(comment, locator),
AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtAsyncFunctionDef(_) => {
handle_leading_function_with_decorators_comment(comment)
}
AnyNodeRef::StmtFunctionDef(_) => handle_leading_function_with_decorators_comment(comment),
AnyNodeRef::StmtClassDef(class_def) => {
handle_leading_class_with_decorators_comment(comment, class_def)
}
@ -178,7 +176,6 @@ fn handle_end_of_line_comment_around_body<'a>(
fn is_first_statement_in_body(statement: AnyNodeRef, has_body: AnyNodeRef) -> bool {
match has_body {
AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { body, orelse, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => {
are_same_optional(statement, body.first())
|| are_same_optional(statement, orelse.first())
@ -208,7 +205,6 @@ fn is_first_statement_in_body(statement: AnyNodeRef, has_body: AnyNodeRef) -> bo
body, ..
})
| AnyNodeRef::StmtFunctionDef(ast::StmtFunctionDef { body, .. })
| AnyNodeRef::StmtAsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. })
| AnyNodeRef::StmtClassDef(ast::StmtClassDef { body, .. }) => {
are_same_optional(statement, body.first())
}
@ -773,9 +769,7 @@ fn handle_module_level_own_line_comment_before_class_or_function_comment<'a>(
// ... where the following is a function or class statement.
if !matches!(
following,
AnyNodeRef::StmtAsyncFunctionDef(_)
| AnyNodeRef::StmtFunctionDef(_)
| AnyNodeRef::StmtClassDef(_)
AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtClassDef(_)
) {
return CommentPlacement::Default(comment);
}
@ -1408,10 +1402,8 @@ where
fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
let body = match node {
AnyNodeRef::StmtFunctionDef(ast::StmtFunctionDef { body, .. })
| AnyNodeRef::StmtAsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. })
| AnyNodeRef::StmtClassDef(ast::StmtClassDef { body, .. })
| AnyNodeRef::StmtWith(ast::StmtWith { body, .. })
| AnyNodeRef::StmtAsyncWith(ast::StmtAsyncWith { body, .. })
| AnyNodeRef::MatchCase(MatchCase { body, .. })
| AnyNodeRef::ExceptHandlerExceptHandler(ast::ExceptHandlerExceptHandler {
body, ..
@ -1424,7 +1416,6 @@ fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
}) => elif_else_clauses.last().map_or(body, |clause| &clause.body),
AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { body, orelse, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => {
if orelse.is_empty() {
body
@ -1477,7 +1468,6 @@ fn last_child_in_body(node: AnyNodeRef) -> Option<AnyNodeRef> {
fn is_first_statement_in_alternate_body(statement: AnyNodeRef, has_body: AnyNodeRef) -> bool {
match has_body {
AnyNodeRef::StmtFor(ast::StmtFor { orelse, .. })
| AnyNodeRef::StmtAsyncFor(ast::StmtAsyncFor { orelse, .. })
| AnyNodeRef::StmtWhile(ast::StmtWhile { orelse, .. }) => {
are_same_optional(statement, orelse.first())
}

1
crates/ruff_python_formatter/src/expression/expr_named_expr.rs
View file

@ -46,7 +46,6 @@ impl NeedsParentheses for ExprNamedExpr {
|| parent.is_with_item()
|| parent.is_stmt_delete()
|| parent.is_stmt_for()
|| parent.is_stmt_async_for()
{
OptionalParentheses::Always
} else {

108
crates/ruff_python_formatter/src/generated.rs
View file

@ -108,42 +108,6 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtFunctionDef {
}
}
impl FormatRule<ast::StmtAsyncFunctionDef, PyFormatContext<'_>>
for crate::statement::stmt_async_function_def::FormatStmtAsyncFunctionDef
{
#[inline]
fn fmt(&self, node: &ast::StmtAsyncFunctionDef, f: &mut PyFormatter) -> FormatResult<()> {
FormatNodeRule::<ast::StmtAsyncFunctionDef>::fmt(self, node, f)
}
}
impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::StmtAsyncFunctionDef {
type Format<'a> = FormatRefWithRule<
'a,
ast::StmtAsyncFunctionDef,
crate::statement::stmt_async_function_def::FormatStmtAsyncFunctionDef,
PyFormatContext<'ast>,
>;
fn format(&self) -> Self::Format<'_> {
FormatRefWithRule::new(
self,
crate::statement::stmt_async_function_def::FormatStmtAsyncFunctionDef::default(),
)
}
}
impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtAsyncFunctionDef {
type Format = FormatOwnedWithRule<
ast::StmtAsyncFunctionDef,
crate::statement::stmt_async_function_def::FormatStmtAsyncFunctionDef,
PyFormatContext<'ast>,
>;
fn into_format(self) -> Self::Format {
FormatOwnedWithRule::new(
self,
crate::statement::stmt_async_function_def::FormatStmtAsyncFunctionDef::default(),
)
}
}
impl FormatRule<ast::StmtClassDef, PyFormatContext<'_>>
for crate::statement::stmt_class_def::FormatStmtClassDef
{
@ -424,42 +388,6 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtFor {
}
}
impl FormatRule<ast::StmtAsyncFor, PyFormatContext<'_>>
for crate::statement::stmt_async_for::FormatStmtAsyncFor
{
#[inline]
fn fmt(&self, node: &ast::StmtAsyncFor, f: &mut PyFormatter) -> FormatResult<()> {
FormatNodeRule::<ast::StmtAsyncFor>::fmt(self, node, f)
}
}
impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::StmtAsyncFor {
type Format<'a> = FormatRefWithRule<
'a,
ast::StmtAsyncFor,
crate::statement::stmt_async_for::FormatStmtAsyncFor,
PyFormatContext<'ast>,
>;
fn format(&self) -> Self::Format<'_> {
FormatRefWithRule::new(
self,
crate::statement::stmt_async_for::FormatStmtAsyncFor::default(),
)
}
}
impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtAsyncFor {
type Format = FormatOwnedWithRule<
ast::StmtAsyncFor,
crate::statement::stmt_async_for::FormatStmtAsyncFor,
PyFormatContext<'ast>,
>;
fn into_format(self) -> Self::Format {
FormatOwnedWithRule::new(
self,
crate::statement::stmt_async_for::FormatStmtAsyncFor::default(),
)
}
}
impl FormatRule<ast::StmtWhile, PyFormatContext<'_>>
for crate::statement::stmt_while::FormatStmtWhile
{
@ -554,42 +482,6 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtWith {
}
}
impl FormatRule<ast::StmtAsyncWith, PyFormatContext<'_>>
for crate::statement::stmt_async_with::FormatStmtAsyncWith
{
#[inline]
fn fmt(&self, node: &ast::StmtAsyncWith, f: &mut PyFormatter) -> FormatResult<()> {
FormatNodeRule::<ast::StmtAsyncWith>::fmt(self, node, f)
}
}
impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::StmtAsyncWith {
type Format<'a> = FormatRefWithRule<
'a,
ast::StmtAsyncWith,
crate::statement::stmt_async_with::FormatStmtAsyncWith,
PyFormatContext<'ast>,
>;
fn format(&self) -> Self::Format<'_> {
FormatRefWithRule::new(
self,
crate::statement::stmt_async_with::FormatStmtAsyncWith::default(),
)
}
}
impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::StmtAsyncWith {
type Format = FormatOwnedWithRule<
ast::StmtAsyncWith,
crate::statement::stmt_async_with::FormatStmtAsyncWith,
PyFormatContext<'ast>,
>;
fn into_format(self) -> Self::Format {
FormatOwnedWithRule::new(
self,
crate::statement::stmt_async_with::FormatStmtAsyncWith::default(),
)
}
}
impl FormatRule<ast::StmtMatch, PyFormatContext<'_>>
for crate::statement::stmt_match::FormatStmtMatch
{

6
crates/ruff_python_formatter/src/statement/mod.rs
View file

@ -5,9 +5,6 @@ use ruff_python_ast::Stmt;
pub(crate) mod stmt_ann_assign;
pub(crate) mod stmt_assert;
pub(crate) mod stmt_assign;
pub(crate) mod stmt_async_for;
pub(crate) mod stmt_async_function_def;
pub(crate) mod stmt_async_with;
pub(crate) mod stmt_aug_assign;
pub(crate) mod stmt_break;
pub(crate) mod stmt_class_def;
@ -40,7 +37,6 @@ impl FormatRule<Stmt, PyFormatContext<'_>> for FormatStmt {
fn fmt(&self, item: &Stmt, f: &mut PyFormatter) -> FormatResult<()> {
match item {
Stmt::FunctionDef(x) => x.format().fmt(f),
Stmt::AsyncFunctionDef(x) => x.format().fmt(f),
Stmt::ClassDef(x) => x.format().fmt(f),
Stmt::Return(x) => x.format().fmt(f),
Stmt::Delete(x) => x.format().fmt(f),
@ -48,11 +44,9 @@ impl FormatRule<Stmt, PyFormatContext<'_>> for FormatStmt {
Stmt::AugAssign(x) => x.format().fmt(f),
Stmt::AnnAssign(x) => x.format().fmt(f),
Stmt::For(x) => x.format().fmt(f),
Stmt::AsyncFor(x) => x.format().fmt(f),
Stmt::While(x) => x.format().fmt(f),
Stmt::If(x) => x.format().fmt(f),
Stmt::With(x) => x.format().fmt(f),
Stmt::AsyncWith(x) => x.format().fmt(f),
Stmt::Match(x) => x.format().fmt(f),
Stmt::Raise(x) => x.format().fmt(f),
Stmt::Try(x) => x.format().fmt(f),

23
crates/ruff_python_formatter/src/statement/stmt_async_for.rs
View file

@ -1,23 +0,0 @@
use crate::prelude::*;
use crate::statement::stmt_for::AnyStatementFor;
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::FormatResult;
use ruff_python_ast::StmtAsyncFor;
#[derive(Default)]
pub struct FormatStmtAsyncFor;
impl FormatNodeRule<StmtAsyncFor> for FormatStmtAsyncFor {
fn fmt_fields(&self, item: &StmtAsyncFor, f: &mut PyFormatter) -> FormatResult<()> {
AnyStatementFor::from(item).fmt(f)
}
fn fmt_dangling_comments(
&self,
_node: &StmtAsyncFor,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}

24
crates/ruff_python_formatter/src/statement/stmt_async_function_def.rs
View file

@ -1,24 +0,0 @@
use ruff_python_ast::StmtAsyncFunctionDef;
use ruff_python_ast::function::AnyFunctionDefinition;
use crate::prelude::*;
use crate::FormatNodeRule;
#[derive(Default)]
pub struct FormatStmtAsyncFunctionDef;
impl FormatNodeRule<StmtAsyncFunctionDef> for FormatStmtAsyncFunctionDef {
fn fmt_fields(&self, item: &StmtAsyncFunctionDef, f: &mut PyFormatter) -> FormatResult<()> {
AnyFunctionDefinition::from(item).format().fmt(f)
}
fn fmt_dangling_comments(
&self,
_node: &StmtAsyncFunctionDef,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled by `AnyFunctionDef`
Ok(())
}
}

22
crates/ruff_python_formatter/src/statement/stmt_async_with.rs
View file

@ -1,22 +0,0 @@
use crate::prelude::*;
use crate::statement::stmt_with::AnyStatementWith;
use crate::FormatNodeRule;
use ruff_python_ast::StmtAsyncWith;
#[derive(Default)]
pub struct FormatStmtAsyncWith;
impl FormatNodeRule<StmtAsyncWith> for FormatStmtAsyncWith {
fn fmt_fields(&self, item: &StmtAsyncWith, f: &mut PyFormatter) -> FormatResult<()> {
AnyStatementWith::from(item).fmt(f)
}
fn fmt_dangling_comments(
&self,
_node: &StmtAsyncWith,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}

106
crates/ruff_python_formatter/src/statement/stmt_for.rs
View file

@ -1,13 +1,12 @@
use ruff_formatter::{format_args, write, Buffer, FormatResult};
use ruff_python_ast::{Expr, Ranged, Stmt, StmtFor};
use crate::comments::{leading_alternate_branch_comments, trailing_comments};
use crate::expression::expr_tuple::TupleParentheses;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::{FormatNodeRule, PyFormatter};
use ruff_formatter::{format_args, write, Buffer, FormatResult};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{Expr, Ranged, Stmt, StmtAsyncFor, StmtFor, Suite};
use ruff_text_size::TextRange;
#[derive(Debug)]
struct ExprTupleWithoutParentheses<'a>(&'a Expr);
@ -27,85 +26,19 @@ impl Format<PyFormatContext<'_>> for ExprTupleWithoutParentheses<'_> {
#[derive(Default)]
pub struct FormatStmtFor;
pub(super) enum AnyStatementFor<'a> {
For(&'a StmtFor),
AsyncFor(&'a StmtAsyncFor),
}
impl<'a> AnyStatementFor<'a> {
const fn is_async(&self) -> bool {
matches!(self, AnyStatementFor::AsyncFor(_))
}
fn target(&self) -> &Expr {
match self {
AnyStatementFor::For(stmt) => &stmt.target,
AnyStatementFor::AsyncFor(stmt) => &stmt.target,
}
}
#[allow(clippy::iter_not_returning_iterator)]
fn iter(&self) -> &Expr {
match self {
AnyStatementFor::For(stmt) => &stmt.iter,
AnyStatementFor::AsyncFor(stmt) => &stmt.iter,
}
}
fn body(&self) -> &Suite {
match self {
AnyStatementFor::For(stmt) => &stmt.body,
AnyStatementFor::AsyncFor(stmt) => &stmt.body,
}
}
fn orelse(&self) -> &Suite {
match self {
AnyStatementFor::For(stmt) => &stmt.orelse,
AnyStatementFor::AsyncFor(stmt) => &stmt.orelse,
}
}
}
impl Ranged for AnyStatementFor<'_> {
fn range(&self) -> TextRange {
match self {
AnyStatementFor::For(stmt) => stmt.range(),
AnyStatementFor::AsyncFor(stmt) => stmt.range(),
}
}
}
impl<'a> From<&'a StmtFor> for AnyStatementFor<'a> {
fn from(value: &'a StmtFor) -> Self {
AnyStatementFor::For(value)
}
}
impl<'a> From<&'a StmtAsyncFor> for AnyStatementFor<'a> {
fn from(value: &'a StmtAsyncFor) -> Self {
AnyStatementFor::AsyncFor(value)
}
}
impl<'a> From<&AnyStatementFor<'a>> for AnyNodeRef<'a> {
fn from(value: &AnyStatementFor<'a>) -> Self {
match value {
AnyStatementFor::For(stmt) => AnyNodeRef::StmtFor(stmt),
AnyStatementFor::AsyncFor(stmt) => AnyNodeRef::StmtAsyncFor(stmt),
}
}
}
impl Format<PyFormatContext<'_>> for AnyStatementFor<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
let target = self.target();
let iter = self.iter();
let body = self.body();
let orelse = self.orelse();
impl FormatNodeRule<StmtFor> for FormatStmtFor {
fn fmt_fields(&self, item: &StmtFor, f: &mut PyFormatter) -> FormatResult<()> {
let StmtFor {
is_async,
target,
iter,
body,
orelse,
range: _,
} = item;
let comments = f.context().comments().clone();
let dangling_comments = comments.dangling_comments(self);
let dangling_comments = comments.dangling_comments(item);
let body_start = body.first().map_or(iter.end(), Stmt::start);
let or_else_comments_start =
dangling_comments.partition_point(|comment| comment.slice().end() < body_start);
@ -116,15 +49,14 @@ impl Format<PyFormatContext<'_>> for AnyStatementFor<'_> {
write!(
f,
[
self.is_async()
.then_some(format_args![text("async"), space()]),
is_async.then_some(format_args![text("async"), space()]),
text("for"),
space(),
ExprTupleWithoutParentheses(target),
space(),
text("in"),
space(),
maybe_parenthesize_expression(iter, self, Parenthesize::IfBreaks),
maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks),
text(":"),
trailing_comments(trailing_condition_comments),
block_indent(&body.format())
@ -153,12 +85,6 @@ impl Format<PyFormatContext<'_>> for AnyStatementFor<'_> {
Ok(())
}
}
impl FormatNodeRule<StmtFor> for FormatStmtFor {
fn fmt_fields(&self, item: &StmtFor, f: &mut PyFormatter) -> FormatResult<()> {
AnyStatementFor::from(item).fmt(f)
}
fn fmt_dangling_comments(&self, _node: &StmtFor, _f: &mut PyFormatter) -> FormatResult<()> {
// Handled in `fmt_fields`

71
crates/ruff_python_formatter/src/statement/stmt_function_def.rs
View file

@ -1,11 +1,8 @@
use ruff_formatter::write;
use ruff_python_ast::{Ranged, StmtFunctionDef};
use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule};
use ruff_python_ast::function::AnyFunctionDefinition;
use ruff_python_trivia::lines_after_ignoring_trivia;
use crate::comments::{leading_comments, trailing_comments};
use crate::expression::parentheses::{optional_parentheses, Parentheses};
use crate::prelude::*;
use crate::statement::suite::SuiteKind;
@ -16,24 +13,6 @@ pub struct FormatStmtFunctionDef;
impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
fn fmt_fields(&self, item: &StmtFunctionDef, f: &mut PyFormatter) -> FormatResult<()> {
AnyFunctionDefinition::from(item).format().fmt(f)
}
fn fmt_dangling_comments(
&self,
_node: &StmtFunctionDef,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled by `AnyFunctionDef`
Ok(())
}
}
#[derive(Default)]
pub struct FormatAnyFunctionDef;
impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFunctionDef {
fn fmt(&self, item: &AnyFunctionDefinition<'_>, f: &mut PyFormatter) -> FormatResult<()> {
let comments = f.context().comments().clone();
let dangling_comments = comments.dangling_comments(item);
@ -43,9 +22,9 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
let (leading_definition_comments, trailing_definition_comments) =
dangling_comments.split_at(trailing_definition_comments_start);
if let Some(last_decorator) = item.decorators().last() {
if let Some(last_decorator) = item.decorator_list.last() {
f.join_with(hard_line_break())
.entries(item.decorators().iter().formatted())
.entries(item.decorator_list.iter().formatted())
.finish()?;
if leading_definition_comments.is_empty() {
@ -69,21 +48,19 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
}
}
if item.is_async() {
if item.is_async {
write!(f, [text("async"), space()])?;
}
let name = item.name();
write!(f, [text("def"), space(), item.name.format()])?;
write!(f, [text("def"), space(), name.format()])?;
if let Some(type_params) = item.type_params() {
if let Some(type_params) = item.type_params.as_ref() {
write!(f, [type_params.format()])?;
}
write!(f, [item.arguments().format()])?;
write!(f, [item.parameters.format()])?;
if let Some(return_annotation) = item.returns() {
if let Some(return_annotation) = item.returns.as_ref() {
write!(
f,
[
@ -102,33 +79,17 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
[
text(":"),
trailing_comments(trailing_definition_comments),
block_indent(&item.body().format().with_options(SuiteKind::Function))
block_indent(&item.body.format().with_options(SuiteKind::Function))
]
)
}
}
impl<'def, 'ast> AsFormat<PyFormatContext<'ast>> for AnyFunctionDefinition<'def> {
type Format<'a> = FormatRefWithRule<
'a,
AnyFunctionDefinition<'def>,
FormatAnyFunctionDef,
PyFormatContext<'ast>,
> where Self: 'a;
fn format(&self) -> Self::Format<'_> {
FormatRefWithRule::new(self, FormatAnyFunctionDef)
}
}
impl<'def, 'ast> IntoFormat<PyFormatContext<'ast>> for AnyFunctionDefinition<'def> {
type Format = FormatOwnedWithRule<
AnyFunctionDefinition<'def>,
FormatAnyFunctionDef,
PyFormatContext<'ast>,
>;
fn into_format(self) -> Self::Format {
FormatOwnedWithRule::new(self, FormatAnyFunctionDef)
fn fmt_dangling_comments(
&self,
_node: &StmtFunctionDef,
_f: &mut PyFormatter,
) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}

106
crates/ruff_python_formatter/src/statement/stmt_with.rs
View file

@ -1,9 +1,7 @@
use ruff_python_ast::{Ranged, StmtAsyncWith, StmtWith, Suite, WithItem};
use ruff_text_size::TextRange;
use ruff_formatter::{format_args, write, FormatError};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{Ranged, StmtWith};
use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
use ruff_text_size::TextRange;
use crate::comments::trailing_comments;
use crate::expression::parentheses::{
@ -12,81 +10,29 @@ use crate::expression::parentheses::{
use crate::prelude::*;
use crate::FormatNodeRule;
pub(super) enum AnyStatementWith<'a> {
With(&'a StmtWith),
AsyncWith(&'a StmtAsyncWith),
}
#[derive(Default)]
pub struct FormatStmtWith;
impl<'a> AnyStatementWith<'a> {
const fn is_async(&self) -> bool {
matches!(self, AnyStatementWith::AsyncWith(_))
}
fn items(&self) -> &[WithItem] {
match self {
AnyStatementWith::With(with) => with.items.as_slice(),
AnyStatementWith::AsyncWith(with) => with.items.as_slice(),
}
}
fn body(&self) -> &Suite {
match self {
AnyStatementWith::With(with) => &with.body,
AnyStatementWith::AsyncWith(with) => &with.body,
}
}
}
impl Ranged for AnyStatementWith<'_> {
fn range(&self) -> TextRange {
match self {
AnyStatementWith::With(with) => with.range(),
AnyStatementWith::AsyncWith(with) => with.range(),
}
}
}
impl<'a> From<&'a StmtWith> for AnyStatementWith<'a> {
fn from(value: &'a StmtWith) -> Self {
AnyStatementWith::With(value)
}
}
impl<'a> From<&'a StmtAsyncWith> for AnyStatementWith<'a> {
fn from(value: &'a StmtAsyncWith) -> Self {
AnyStatementWith::AsyncWith(value)
}
}
impl<'a> From<&AnyStatementWith<'a>> for AnyNodeRef<'a> {
fn from(value: &AnyStatementWith<'a>) -> Self {
match value {
AnyStatementWith::With(with) => AnyNodeRef::StmtWith(with),
AnyStatementWith::AsyncWith(with) => AnyNodeRef::StmtAsyncWith(with),
}
}
}
impl Format<PyFormatContext<'_>> for AnyStatementWith<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
impl FormatNodeRule<StmtWith> for FormatStmtWith {
fn fmt_fields(&self, item: &StmtWith, f: &mut PyFormatter) -> FormatResult<()> {
let comments = f.context().comments().clone();
let dangling_comments = comments.dangling_comments(self);
let dangling_comments = comments.dangling_comments(item);
write!(
f,
[
self.is_async()
item.is_async
.then_some(format_args![text("async"), space()]),
text("with"),
space()
]
)?;
if are_with_items_parenthesized(self, f.context())? {
if are_with_items_parenthesized(item, f.context())? {
optional_parentheses(&format_with(|f| {
let mut joiner = f.join_comma_separated(self.body().first().unwrap().start());
let mut joiner = f.join_comma_separated(item.body.first().unwrap().start());
for item in self.items() {
for item in &item.items {
joiner.entry_with_line_separator(
item,
&item.format(),
@ -98,7 +44,7 @@ impl Format<PyFormatContext<'_>> for AnyStatementWith<'_> {
.fmt(f)?;
} else {
f.join_with(format_args![text(","), space()])
.entries(self.items().iter().formatted())
.entries(item.items.iter().formatted())
.finish()?;
}
@ -107,18 +53,20 @@ impl Format<PyFormatContext<'_>> for AnyStatementWith<'_> {
[
text(":"),
trailing_comments(dangling_comments),
block_indent(&self.body().format())
block_indent(&item.body.format())
]
)
}
fn fmt_dangling_comments(&self, _node: &StmtWith, _f: &mut PyFormatter) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}
fn are_with_items_parenthesized(
with: &AnyStatementWith,
context: &PyFormatContext,
) -> FormatResult<bool> {
fn are_with_items_parenthesized(with: &StmtWith, context: &PyFormatContext) -> FormatResult<bool> {
let first_with_item = with
.items()
.items
.first()
.ok_or(FormatError::syntax_error("Expected at least one with item"))?;
let before_first_with_item = TextRange::new(with.start(), first_with_item.start());
@ -145,17 +93,3 @@ fn are_with_items_parenthesized(
None => Ok(false),
}
}
#[derive(Default)]
pub struct FormatStmtWith;
impl FormatNodeRule<StmtWith> for FormatStmtWith {
fn fmt_fields(&self, item: &StmtWith, f: &mut PyFormatter) -> FormatResult<()> {
AnyStatementWith::from(item).fmt(f)
}
fn fmt_dangling_comments(&self, _node: &StmtWith, _f: &mut PyFormatter) -> FormatResult<()> {
// Handled in `fmt_fields`
Ok(())
}
}

5
crates/ruff_python_formatter/src/statement/suite.rs
View file

@ -210,10 +210,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
/// Returns `true` if a [`Stmt`] is a class or function definition.
const fn is_class_or_function_definition(stmt: &Stmt) -> bool {
matches!(
stmt,
Stmt::FunctionDef(_) | Stmt::AsyncFunctionDef(_) | Stmt::ClassDef(_)
)
matches!(stmt, Stmt::FunctionDef(_) | Stmt::ClassDef(_))
}
/// Returns `true` if a [`Stmt`] is an import.