Enable function2 test (#3083)

crates/ruff_python_formatter/src/format/stmt.rs

@@ -322,7 +322,7 @@ fn format_for(
     target: &Expr,
     iter: &Expr,
     body: &[Stmt],
-    _orelse: &[Stmt],
+    orelse: &[Stmt],
     _type_comment: Option<&str>,
 ) -> FormatResult<()> {
     write!(
@@ -338,7 +338,11 @@ fn format_for(
             text(":"),
             block_indent(&block(body))
         ]
-    )
+    )?;
+    if !orelse.is_empty() {
+        write!(f, [text("else:"), block_indent(&block(orelse))])?;
+    }
+    Ok(())
 }
 
 fn format_while(

crates/ruff_python_formatter/src/lib.rs

@@ -76,6 +76,8 @@ mod tests {
     #[test_case(Path::new("simple_cases/tricky_unicode_symbols.py"); "tricky_unicode_symbols")]
     // Passing except that `1, 2, 3,` should be `(1, 2, 3)`.
     #[test_case(Path::new("simple_cases/tupleassign.py"); "tupleassign")]
+    // Passing except that `CliRunner().invoke(...)` arguments are improperly wrapped.
+    #[test_case(Path::new("simple_cases/function2.py"); "function2")]
     fn passing(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.display());
         let content = std::fs::read_to_string(test_resource_path(
@@ -110,7 +112,6 @@ mod tests {
     // inappropriately associated with the if statement rather than the line it's on.
     #[test_case(Path::new("simple_cases/comments.py"); "comments")]
     #[test_case(Path::new("simple_cases/function.py"); "function")]
-    #[test_case(Path::new("simple_cases/function2.py"); "function2")]
     #[test_case(Path::new("simple_cases/function_trailing_comma.py"); "function_trailing_comma")]
     #[test_case(Path::new("simple_cases/composition.py"); "composition")]
     fn failing(path: &Path) -> Result<()> {

crates/ruff_python_formatter/src/newlines.rs

@@ -2,10 +2,10 @@ use rustpython_parser::ast::Constant;
 
 use crate::core::visitor;
 use crate::core::visitor::Visitor;
-use crate::cst::{Expr, ExprKind, Stmt, StmtKind};
+use crate::cst::{ExcepthandlerKind, Expr, ExprKind, Stmt, StmtKind};
 use crate::trivia::{Relationship, Trivia, TriviaKind};
 
-#[derive(Copy, Clone)]
+#[derive(Debug, Copy, Clone)]
 enum Depth {
     TopLevel,
     Nested,
@@ -20,7 +20,7 @@ impl Depth {
     }
 }
 
-#[derive(Copy, Clone)]
+#[derive(Debug, Copy, Clone)]
 enum Scope {
     Module,
     Class,
@@ -35,6 +35,7 @@ enum Trailer {
     Import,
     Docstring,
     Generic,
+    CompoundStatement,
 }
 
 struct NewlineNormalizer {
@@ -60,10 +61,18 @@ impl<'a> Visitor<'a> for NewlineNormalizer {
             }
         });
 
-        if matches!(self.trailer, Trailer::None) {
+        if matches!(self.trailer, Trailer::None)
-            // If this is the first statement in the block, remove any leading empty lines.
+            || (matches!(self.trailer, Trailer::CompoundStatement)
-            // TODO(charlie): If we have a function or class definition within a non-scoped block,
+                && !matches!(
-            // like an if-statement, retain a line before and after.
+                    stmt.node,
+                    StmtKind::FunctionDef { .. }
+                        | StmtKind::AsyncFunctionDef { .. }
+                        | StmtKind::ClassDef { .. }
+                ))
+        {
+            // If this is the first statement in the block, remove any leading empty lines, with the
+            // exception being functions and classes defined within compound statements (e.g., as
+            // the first statement in an `if` body).
             let mut seen_non_empty = false;
             stmt.trivia.retain(|c| {
                 if seen_non_empty {
@@ -145,12 +154,138 @@ impl<'a> Visitor<'a> for NewlineNormalizer {
             }
         }
 
-        self.trailer = match &stmt.node {
+        let prev_scope = self.scope;
-            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+        let prev_depth = self.depth;
-                Trailer::FunctionDef
+
+        match &mut stmt.node {
+            StmtKind::FunctionDef { body, .. } | StmtKind::AsyncFunctionDef { body, .. } => {
+                self.depth = Depth::Nested;
+                self.scope = Scope::Function;
+                self.trailer = Trailer::None;
+                self.visit_body(body);
+                self.trailer = Trailer::FunctionDef;
             }
-            // TODO(charlie): This needs to be the first statement in a class or function.
+            StmtKind::ClassDef { body, .. } => {
-            StmtKind::Expr { value, .. } => {
+                self.depth = Depth::Nested;
+                self.scope = Scope::Class;
+                self.trailer = Trailer::None;
+                self.visit_body(body);
+                self.trailer = Trailer::ClassDef;
+            }
+            StmtKind::While { body, orelse, .. }
+            | StmtKind::For { body, orelse, .. }
+            | StmtKind::AsyncFor { body, orelse, .. } => {
+                self.depth = Depth::Nested;
+                self.trailer = Trailer::CompoundStatement;
+                self.visit_body(body);
+
+                if !orelse.is_empty() {
+                    // If the previous body ended with a function or class definition, we need to
+                    // insert an empty line before the else block. Since the `else` itself isn't
+                    // a statement, we need to insert it into the last statement of the body.
+                    if matches!(self.trailer, Trailer::ClassDef | Trailer::FunctionDef) {
+                        let stmt = body.last_mut().unwrap();
+                        stmt.trivia.push(Trivia {
+                            kind: TriviaKind::EmptyLine,
+                            relationship: Relationship::Trailing,
+                        });
+                    }
+
+                    self.depth = Depth::Nested;
+                    self.trailer = Trailer::CompoundStatement;
+                    self.visit_body(orelse);
+                }
+            }
+            StmtKind::If { body, orelse, .. } => {
+                self.depth = Depth::Nested;
+                self.trailer = Trailer::CompoundStatement;
+                self.visit_body(body);
+
+                if !orelse.is_empty() {
+                    if matches!(self.trailer, Trailer::ClassDef | Trailer::FunctionDef) {
+                        let stmt = body.last_mut().unwrap();
+                        stmt.trivia.push(Trivia {
+                            kind: TriviaKind::EmptyLine,
+                            relationship: Relationship::Trailing,
+                        });
+                    }
+
+                    self.depth = Depth::Nested;
+                    self.trailer = Trailer::CompoundStatement;
+                    self.visit_body(orelse);
+                }
+            }
+            StmtKind::With { body, .. } | StmtKind::AsyncWith { body, .. } => {
+                self.depth = Depth::Nested;
+                self.trailer = Trailer::CompoundStatement;
+                self.visit_body(body);
+            }
+            // StmtKind::Match { .. } => {}
+            StmtKind::Try {
+                body,
+                handlers,
+                orelse,
+                finalbody,
+            } => {
+                self.depth = Depth::Nested;
+                self.trailer = Trailer::CompoundStatement;
+                self.visit_body(body);
+                let mut last = body.last_mut();
+
+                for handler in handlers {
+                    if matches!(self.trailer, Trailer::ClassDef | Trailer::FunctionDef) {
+                        if let Some(stmt) = last.as_mut() {
+                            stmt.trivia.push(Trivia {
+                                kind: TriviaKind::EmptyLine,
+                                relationship: Relationship::Trailing,
+                            });
+                        }
+                    }
+
+                    self.depth = Depth::Nested;
+                    self.trailer = Trailer::CompoundStatement;
+                    let ExcepthandlerKind::ExceptHandler { body, .. } = &mut handler.node;
+                    self.visit_body(body);
+                    last = body.last_mut();
+                }
+
+                if !orelse.is_empty() {
+                    if matches!(self.trailer, Trailer::ClassDef | Trailer::FunctionDef) {
+                        if let Some(stmt) = last.as_mut() {
+                            stmt.trivia.push(Trivia {
+                                kind: TriviaKind::EmptyLine,
+                                relationship: Relationship::Trailing,
+                            });
+                        }
+                    }
+
+                    self.depth = Depth::Nested;
+                    self.trailer = Trailer::CompoundStatement;
+                    self.visit_body(orelse);
+                    last = body.last_mut();
+                }
+
+                if !finalbody.is_empty() {
+                    if matches!(self.trailer, Trailer::ClassDef | Trailer::FunctionDef) {
+                        if let Some(stmt) = last.as_mut() {
+                            stmt.trivia.push(Trivia {
+                                kind: TriviaKind::EmptyLine,
+                                relationship: Relationship::Trailing,
+                            });
+                        }
+                    }
+
+                    self.depth = Depth::Nested;
+                    self.trailer = Trailer::CompoundStatement;
+                    self.visit_body(finalbody);
+                }
+            }
+            _ => {
+                self.trailer = match &stmt.node {
+                    StmtKind::Expr { value, .. }
+                        if matches!(self.scope, Scope::Class | Scope::Function)
+                            && matches!(self.trailer, Trailer::None) =>
+                    {
                         if let ExprKind::Constant {
                             value: Constant::Str(..),
                             ..
@@ -161,41 +296,18 @@ impl<'a> Visitor<'a> for NewlineNormalizer {
                             Trailer::Generic
                         }
                     }
-            StmtKind::ClassDef { .. } => Trailer::ClassDef,
                     StmtKind::Import { .. } | StmtKind::ImportFrom { .. } => Trailer::Import,
                     _ => Trailer::Generic,
                 };
-
-        let prev_scope = self.scope;
-        self.scope = match &stmt.node {
-            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => Scope::Function,
-            StmtKind::ClassDef { .. } => Scope::Class,
-            _ => prev_scope,
-        };
-
                 visitor::walk_stmt(self, stmt);
+            }
+        }
 
+        self.depth = prev_depth;
         self.scope = prev_scope;
     }
 
-    fn visit_expr(&mut self, expr: &'a mut Expr) {
+    fn visit_expr(&mut self, _expr: &'a mut Expr) {}
-        expr.trivia
-            .retain(|c| !matches!(c.kind, TriviaKind::EmptyLine));
-        visitor::walk_expr(self, expr);
-    }
-
-    fn visit_body(&mut self, body: &'a mut [Stmt]) {
-        let prev_depth = self.depth;
-        let prev_trailer = self.trailer;
-
-        self.depth = Depth::Nested;
-        self.trailer = Trailer::None;
-
-        visitor::walk_body(self, body);
-
-        self.trailer = prev_trailer;
-        self.depth = prev_depth;
-    }
 }
 
 pub fn normalize_newlines(python_cst: &mut [Stmt]) {

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__simple_cases__function2.py.snap

@@ -1,7 +1,6 @@
 ---
-source: src/source_code/mod.rs
+source: crates/ruff_python_formatter/src/lib.rs
-assertion_line: 0
+expression: adjust_quotes(formatted.print()?.as_code())
-expression: formatted
 ---
 def f(
     a,
@@ -10,7 +9,8 @@ def f(
     with cache_dir():
         if something:
             result = CliRunner().invoke(
-                black.main, [str(src1), str(src2), "--diff", "--check"]
+                black.main,
+                [str(src1), str(src2), "--diff", "--check"],
             )
     limited.append(-limited.pop())  # negate top
     return A(