Format AsyncFor (#5808)

crates/ruff_python_formatter/src/statement/stmt_async_for.rs

@@ -1,5 +1,7 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::prelude::*;
+use crate::statement::stmt_for::AnyStatementFor;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::FormatResult;
 use rustpython_parser::ast::StmtAsyncFor;
 
 #[derive(Default)]
@@ -7,6 +9,15 @@ pub struct FormatStmtAsyncFor;
 
 impl FormatNodeRule<StmtAsyncFor> for FormatStmtAsyncFor {
     fn fmt_fields(&self, item: &StmtAsyncFor, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        AnyStatementFor::from(item).fmt(f)
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &StmtAsyncFor,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_for.rs

@@ -4,9 +4,10 @@ use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
 use crate::{FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use ruff_python_ast::node::AstNode;
-use rustpython_parser::ast::{Expr, Ranged, Stmt, StmtFor};
+use ruff_formatter::{format_args, write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
+use ruff_text_size::TextRange;
+use rustpython_parser::ast::{Expr, Ranged, Stmt, StmtAsyncFor, StmtFor, Suite};
 
 #[derive(Debug)]
 struct ExprTupleWithoutParentheses<'a>(&'a Expr);
@@ -26,19 +27,85 @@ impl Format<PyFormatContext<'_>> for ExprTupleWithoutParentheses<'_> {
 #[derive(Default)]
 pub struct FormatStmtFor;
 
-impl FormatNodeRule<StmtFor> for FormatStmtFor {
-    fn fmt_fields(&self, item: &StmtFor, f: &mut PyFormatter) -> FormatResult<()> {
-        let StmtFor {
-            range: _,
-            target,
-            iter,
-            body,
-            orelse,
-            type_comment: _,
-        } = item;
+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 Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let target = self.target();
+        let iter = self.iter();
+        let body = self.body();
+        let orelse = self.orelse();
 
         let comments = f.context().comments().clone();
-        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+        let dangling_comments = comments.dangling_comments(self);
         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);
@@ -49,13 +116,15 @@ impl FormatNodeRule<StmtFor> for FormatStmtFor {
         write!(
             f,
             [
+                self.is_async()
+                    .then_some(format_args![text("async"), space()]),
                 text("for"),
                 space(),
-                ExprTupleWithoutParentheses(target.as_ref()),
+                ExprTupleWithoutParentheses(target),
                 space(),
                 text("in"),
                 space(),
-                maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks),
+                maybe_parenthesize_expression(iter, self, Parenthesize::IfBreaks),
                 text(":"),
                 trailing_comments(trailing_condition_comments),
                 block_indent(&body.format())
@@ -84,6 +153,12 @@ impl FormatNodeRule<StmtFor> for FormatStmtFor {
 
         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`