diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/assign.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/assign.py index 52b143edb9..a0317eac9a 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/assign.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/assign.py @@ -8,3 +8,11 @@ a2 = ( # Break the last element a = asdf = fjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfal = 1 + +aa = [ + bakjdshflkjahdslkfjlasfdahjlfds +] = dddd = ddd = fkjaödkjaföjfahlfdalfhaöfaöfhaöfha = g = [3] + +aaaa = ( # trailing + # comment + bbbbb) = cccccccccccccccc = 3 diff --git a/crates/ruff_python_formatter/src/expression/mod.rs b/crates/ruff_python_formatter/src/expression/mod.rs index 121610b81c..617fd61c41 100644 --- a/crates/ruff_python_formatter/src/expression/mod.rs +++ b/crates/ruff_python_formatter/src/expression/mod.rs @@ -443,6 +443,10 @@ impl<'input> PreorderVisitor<'input> for CanOmitOptionalParenthesesVisitor<'inpu } fn has_parentheses(expr: &Expr, source: &str) -> bool { + has_own_parentheses(expr) || is_expression_parenthesized(AnyNodeRef::from(expr), source) +} + +pub(crate) const fn has_own_parentheses(expr: &Expr) -> bool { matches!( expr, Expr::Dict(_) @@ -454,7 +458,7 @@ fn has_parentheses(expr: &Expr, source: &str) -> bool { | Expr::DictComp(_) | Expr::Call(_) | Expr::Subscript(_) - ) || is_expression_parenthesized(AnyNodeRef::from(expr), source) + ) } #[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)] diff --git a/crates/ruff_python_formatter/src/statement/stmt_assign.rs b/crates/ruff_python_formatter/src/statement/stmt_assign.rs index 36a53891b5..55ff07d949 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_assign.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_assign.rs @@ -1,16 +1,13 @@ -use rustpython_parser::ast::StmtAssign; +use rustpython_parser::ast::{Expr, StmtAssign}; -use ruff_formatter::write; +use ruff_formatter::{format_args, write, FormatError}; -use crate::expression::maybe_parenthesize_expression; -use crate::expression::parentheses::Parenthesize; +use crate::context::NodeLevel; +use crate::expression::parentheses::{Parentheses, Parenthesize}; +use crate::expression::{has_own_parentheses, maybe_parenthesize_expression}; use crate::prelude::*; use crate::FormatNodeRule; -// Note: This currently does wrap but not the black way so the types below likely need to be -// replaced entirely -// - #[derive(Default)] pub struct FormatStmtAssign; @@ -23,9 +20,18 @@ impl FormatNodeRule for FormatStmtAssign { type_comment: _, } = item; - for target in targets { - write!(f, [target.format(), space(), text("="), space()])?; - } + let (first, rest) = targets.split_first().ok_or(FormatError::SyntaxError)?; + + write!( + f, + [ + first.format(), + space(), + text("="), + space(), + FormatTargets { targets: rest } + ] + )?; write!( f, @@ -37,3 +43,58 @@ impl FormatNodeRule for FormatStmtAssign { ) } } + +struct FormatTargets<'a> { + targets: &'a [Expr], +} + +impl Format> for FormatTargets<'_> { + fn fmt(&self, f: &mut Formatter>) -> FormatResult<()> { + if let Some((first, rest)) = self.targets.split_first() { + let can_omit_parentheses = has_own_parentheses(first); + + let group_id = if can_omit_parentheses { + Some(f.group_id("assignment_parentheses")) + } else { + None + }; + + let saved_level = f.context().node_level(); + f.context_mut() + .set_node_level(NodeLevel::Expression(group_id)); + + let format_first = format_with(|f: &mut PyFormatter| { + let result = if can_omit_parentheses { + first.format().with_options(Parentheses::Never).fmt(f) + } else { + write!( + f, + [ + if_group_breaks(&text("(")), + soft_block_indent(&first.format().with_options(Parentheses::Never)), + if_group_breaks(&text(")")) + ] + ) + }; + + f.context_mut().set_node_level(saved_level); + + result + }); + + write!( + f, + [group(&format_args![ + format_first, + space(), + text("="), + space(), + FormatTargets { targets: rest } + ]) + .with_group_id(group_id)] + ) + } else { + Ok(()) + } + } +} diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__assign.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__assign.py.snap index efcd0cca34..4134d2b91e 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__assign.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__assign.py.snap @@ -14,18 +14,43 @@ a2 = ( # Break the last element a = asdf = fjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfal = 1 + +aa = [ + bakjdshflkjahdslkfjlasfdahjlfds +] = dddd = ddd = fkjaödkjaföjfahlfdalfhaöfaöfhaöfha = g = [3] + +aaaa = ( # trailing + # comment + bbbbb) = cccccccccccccccc = 3 ``` ## Output ```py # break left hand side -a1akjdshflkjahdslkfjlasfdahjlfds = bakjdshflkjahdslkfjlasfdahjlfds = cakjdshflkjahdslkfjlasfdahjlfds = kjaödkjaföjfahlfdalfhaöfaöfhaöfha = fkjaödkjaföjfahlfdalfhaöfaöfhaöfha = g = 3 +a1akjdshflkjahdslkfjlasfdahjlfds = ( + bakjdshflkjahdslkfjlasfdahjlfds +) = ( + cakjdshflkjahdslkfjlasfdahjlfds +) = kjaödkjaföjfahlfdalfhaöfaöfhaöfha = fkjaödkjaföjfahlfdalfhaöfaöfhaöfha = g = 3 # join left hand side -a2 = (b2) = 2 +a2 = b2 = 2 # Break the last element -a = asdf = fjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfal = 1 +a = ( + asdf +) = ( + fjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfalflaflapamsakjsdhflakjdslfjhalsdljfal +) = 1 + +aa = [ + bakjdshflkjahdslkfjlasfdahjlfds +] = dddd = ddd = fkjaödkjaföjfahlfdalfhaöfaöfhaöfha = g = [3] + +aaaa = ( # trailing + # comment + bbbbb +) = cccccccccccccccc = 3 ```