From a6eb60cdd58b188e39a06873bb6355faccbc02c4 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 20 Feb 2023 23:37:50 -0500 Subject: [PATCH] Enable `function2` test (#3083) --- .../ruff_python_formatter/src/format/stmt.rs | 8 +- crates/ruff_python_formatter/src/lib.rs | 3 +- crates/ruff_python_formatter/src/newlines.rs | 210 ++++++++++++++---- ...r__tests__simple_cases__function2.py.snap} | 8 +- 4 files changed, 173 insertions(+), 56 deletions(-) rename crates/ruff_python_formatter/src/snapshots/{expect/ruff_python_formatter__tests__simple_cases__function2.py.snap.expect => ruff_python_formatter__tests__simple_cases__function2.py.snap} (85%) diff --git a/crates/ruff_python_formatter/src/format/stmt.rs b/crates/ruff_python_formatter/src/format/stmt.rs index a8742e72b5..6bba65f14b 100644 --- a/crates/ruff_python_formatter/src/format/stmt.rs +++ b/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( diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs index 2ccbe6186c..2b416a45cd 100644 --- a/crates/ruff_python_formatter/src/lib.rs +++ b/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<()> { diff --git a/crates/ruff_python_formatter/src/newlines.rs b/crates/ruff_python_formatter/src/newlines.rs index eae8f4b71b..7f8d48425d 100644 --- a/crates/ruff_python_formatter/src/newlines.rs +++ b/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 this is the first statement in the block, remove any leading empty lines. - // TODO(charlie): If we have a function or class definition within a non-scoped block, - // like an if-statement, retain a line before and after. + if matches!(self.trailer, Trailer::None) + || (matches!(self.trailer, Trailer::CompoundStatement) + && !matches!( + 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,57 +154,160 @@ impl<'a> Visitor<'a> for NewlineNormalizer { } } - self.trailer = match &stmt.node { - StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => { - Trailer::FunctionDef + let prev_scope = self.scope; + let prev_depth = self.depth; + + 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::Expr { value, .. } => { - if let ExprKind::Constant { - value: Constant::Str(..), - .. - } = &value.node - { - Trailer::Docstring - } else { - Trailer::Generic + StmtKind::ClassDef { body, .. } => { + 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::ClassDef { .. } => Trailer::ClassDef, - StmtKind::Import { .. } | StmtKind::ImportFrom { .. } => Trailer::Import, - _ => Trailer::Generic, - }; + StmtKind::If { body, orelse, .. } => { + self.depth = Depth::Nested; + self.trailer = Trailer::CompoundStatement; + self.visit_body(body); - let prev_scope = self.scope; - self.scope = match &stmt.node { - StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => Scope::Function, - StmtKind::ClassDef { .. } => Scope::Class, - _ => prev_scope, - }; + 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, + }); + } - visitor::walk_stmt(self, stmt); + 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(..), + .. + } = &value.node + { + Trailer::Docstring + } else { + Trailer::Generic + } + } + StmtKind::Import { .. } | StmtKind::ImportFrom { .. } => Trailer::Import, + _ => Trailer::Generic, + }; + visitor::walk_stmt(self, stmt); + } + } + + self.depth = prev_depth; self.scope = prev_scope; } - 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; - } + fn visit_expr(&mut self, _expr: &'a mut Expr) {} } pub fn normalize_newlines(python_cst: &mut [Stmt]) { diff --git a/crates/ruff_python_formatter/src/snapshots/expect/ruff_python_formatter__tests__simple_cases__function2.py.snap.expect b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__simple_cases__function2.py.snap similarity index 85% rename from crates/ruff_python_formatter/src/snapshots/expect/ruff_python_formatter__tests__simple_cases__function2.py.snap.expect rename to crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__simple_cases__function2.py.snap index c5249d27ba..559659fc4f 100644 --- a/crates/ruff_python_formatter/src/snapshots/expect/ruff_python_formatter__tests__simple_cases__function2.py.snap.expect +++ b/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 -assertion_line: 0 -expression: formatted +source: crates/ruff_python_formatter/src/lib.rs +expression: adjust_quotes(formatted.print()?.as_code()) --- 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(