Track magic comma

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -2,9 +2,10 @@ use crate::context::NodeLevel;
 use crate::prelude::*;
 use crate::trivia::{first_non_trivia_token, first_non_trivia_token_rev, Token, TokenKind};
 use ruff_formatter::prelude::tag::Condition;
-use ruff_formatter::{format_args, Argument, Arguments};
+use ruff_formatter::{format_args, Argument, Arguments, FormatContext};
 use ruff_python_ast::expression::ExpressionRef;
-use rustpython_parser::ast::Ranged;
+use rustpython_parser::ast;
+use rustpython_parser::ast::{Comprehension, Expr, Ranged};
 
 pub(crate) trait NeedsParentheses {
     fn needs_parentheses(
@@ -46,12 +47,225 @@ pub(super) fn default_expression_needs_parentheses(
 fn can_omit_optional_parentheses(expr: ExpressionRef, context: &PyFormatContext) -> bool {
     if context.comments().has_leading_comments(expr) {
         false
+    } else if context.options().magic_trailing_comma().is_respect()
+        && has_magic_comma(expr, context.source())
+    {
+        false
     } else {
         true
     }
 }
 
-// fn has_magic_comma(expr: AnyNodeRef)
+fn has_magic_comma<'a, E: Into<ExpressionRef<'a>>>(expression: E, source: &str) -> bool {
+    fn any_expr_with_magic_comma(values: &[Expr], source: &str) -> bool {
+        if let Some(last) = values.last() {
+            if matches!(
+                first_non_trivia_token(last.end(), source),
+                Some(Token {
+                    kind: TokenKind::Comma,
+                    ..
+                })
+            ) {
+                return true;
+            }
+        }
+
+        values
+            .iter()
+            .any(|expression| has_magic_comma(expression, source))
+    }
+
+    fn comprehension_has_magic_comma(comprehension: &Comprehension, source: &str) -> bool {
+        let Comprehension {
+            range: _,
+            target,
+            iter,
+            ifs,
+            is_async: _,
+        } = comprehension;
+
+        has_magic_comma(target, source)
+            || has_magic_comma(iter, source)
+            || any_expr_with_magic_comma(&ifs, source)
+    }
+
+    fn has_magic_comma_inner(expr: ExpressionRef, source: &str) -> bool {
+        match expr {
+            ExpressionRef::BoolOp(ast::ExprBoolOp {
+                range: _,
+                op: _,
+                values,
+            }) => values
+                .iter()
+                .any(|expression| has_magic_comma(expression, source)),
+            ExpressionRef::NamedExpr(ast::ExprNamedExpr {
+                value,
+                target,
+                range: _,
+            }) => has_magic_comma(value, source) || has_magic_comma(target, source),
+            ExpressionRef::BinOp(ast::ExprBinOp {
+                range: _,
+                left,
+                op: _,
+                right,
+            }) => has_magic_comma(left, source) || has_magic_comma(right, source),
+            ExpressionRef::UnaryOp(ast::ExprUnaryOp {
+                range: _,
+                op: _,
+                operand,
+            }) => has_magic_comma(operand, source),
+            ExpressionRef::Lambda(ast::ExprLambda {
+                range: _,
+                args: _,
+                body: _,
+            }) => {
+                // Revisit what the behavior for lambda is
+                false
+            }
+            ExpressionRef::IfExp(ast::ExprIfExp {
+                range: _,
+                test,
+                body,
+                orelse,
+            }) => {
+                has_magic_comma(test, source)
+                    || has_magic_comma(body, source)
+                    || has_magic_comma(orelse, source)
+            }
+            ExpressionRef::Dict(ast::ExprDict {
+                keys: _,
+                values,
+                range: _,
+            })
+            | ExpressionRef::Set(ast::ExprSet {
+                range: _,
+                elts: values,
+            })
+            | ExpressionRef::List(ast::ExprList {
+                range: _,
+                elts: values,
+                ctx: _,
+            }) => any_expr_with_magic_comma(&values, source),
+            ExpressionRef::Tuple(ast::ExprTuple {
+                range: _,
+                elts,
+                ctx: _,
+            }) => {
+                if elts.len() > 1 {
+                    any_expr_with_magic_comma(elts, source)
+                } else if let Some(expression) = elts.last() {
+                    has_magic_comma(expression, source)
+                } else {
+                    false
+                }
+            }
+            ExpressionRef::Call(ast::ExprCall {
+                range: _,
+                func,
+                args,
+                keywords: _,
+            }) => has_magic_comma(func, source) || any_expr_with_magic_comma(&args, source),
+            ExpressionRef::ListComp(ast::ExprListComp {
+                range: _,
+                elt,
+                generators,
+            })
+            | ExpressionRef::SetComp(ast::ExprSetComp {
+                range: _,
+                elt,
+                generators,
+            })
+            | ExpressionRef::GeneratorExp(ast::ExprGeneratorExp {
+                range: _,
+                elt,
+                generators,
+            }) => {
+                false
+                // has_magic_comma(elt, source)
+                //     || generators
+                //         .iter()
+                //         .any(|generator| comprehension_has_magic_comma(generator, source))
+            }
+            ExpressionRef::DictComp(ast::ExprDictComp {
+                range: _,
+                key,
+                value,
+                generators,
+            }) => {
+                has_magic_comma(key, source)
+                    || has_magic_comma(value, source)
+                    || generators
+                        .iter()
+                        .any(|generator| comprehension_has_magic_comma(generator, source))
+            }
+            ExpressionRef::Await(ast::ExprAwait { range: _, value })
+            | ExpressionRef::YieldFrom(ast::ExprYieldFrom { range: _, value })
+            | ExpressionRef::Attribute(ast::ExprAttribute {
+                range: _,
+                value,
+                attr: _,
+                ctx: _,
+            })
+            | ExpressionRef::Starred(ast::ExprStarred {
+                range: _,
+                value,
+                ctx: _,
+            }) => has_magic_comma(value, source),
+
+            ExpressionRef::Yield(ast::ExprYield { range: _, value }) => value
+                .as_ref()
+                .map_or(false, |value| has_magic_comma(value, source)),
+
+            ExpressionRef::Compare(ast::ExprCompare {
+                range: _,
+                left,
+                ops: _,
+                comparators,
+            }) => any_expr_with_magic_comma(comparators, source) || has_magic_comma(left, source),
+            ExpressionRef::FormattedValue(ast::ExprFormattedValue {
+                range: _,
+                value,
+                conversion: _,
+                format_spec,
+            }) => {
+                has_magic_comma(value, source)
+                    || format_spec
+                        .as_ref()
+                        .map_or(false, |spec| has_magic_comma(spec, source))
+            }
+            ExpressionRef::JoinedStr(ast::ExprJoinedStr { range: _, values }) => {
+                values.iter().any(|value| has_magic_comma(value, source))
+            }
+
+            ExpressionRef::Subscript(ast::ExprSubscript {
+                range: _,
+                value,
+                slice,
+                ctx: _,
+            }) => has_magic_comma(value, source) || has_magic_comma(slice, source),
+
+            ExpressionRef::Slice(ast::ExprSlice {
+                range: _,
+                lower,
+                upper,
+                step,
+            }) => {
+                lower
+                    .as_ref()
+                    .map_or(false, |lower| has_magic_comma(lower, source))
+                    || upper
+                        .as_ref()
+                        .map_or(false, |upper| has_magic_comma(upper, source))
+                    || step
+                        .as_ref()
+                        .map_or(false, |step| has_magic_comma(step, source))
+            }
+            ExpressionRef::Constant(_) | ExpressionRef::Name(_) => false,
+        }
+    }
+
+    has_magic_comma_inner(expression.into(), source)
+}
 
 /// Configures if the expression should be parenthesized.
 #[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]