implement preview layout

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -10,6 +10,7 @@ use crate::expression::parentheses::{
     NeedsParentheses, OptionalParentheses, Parentheses, is_expression_parenthesized,
 };
 use crate::prelude::*;
+use crate::preview::is_fluent_layout_split_first_call_enabled;
 
 #[derive(Default)]
 pub struct FormatExprAttribute {
@@ -47,20 +48,26 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
                     )
                 };
 
-            if call_chain_layout == CallChainLayout::Fluent {
+            if matches!(call_chain_layout, CallChainLayout::Fluent(_)) {
                 if parenthesize_value {
                     // Don't propagate the call chain layout.
                     value.format().with_options(Parentheses::Always).fmt(f)?;
                 } else {
                     match value.as_ref() {
                         Expr::Attribute(expr) => {
-                            expr.format().with_options(call_chain_layout).fmt(f)?;
+                            expr.format()
+                                .with_options(call_chain_layout.after_attribute())
+                                .fmt(f)?;
                         }
                         Expr::Call(expr) => {
-                            expr.format().with_options(call_chain_layout).fmt(f)?;
+                            expr.format()
+                                .with_options(call_chain_layout.after_attribute())
+                                .fmt(f)?;
                         }
                         Expr::Subscript(expr) => {
-                            expr.format().with_options(call_chain_layout).fmt(f)?;
+                            expr.format()
+                                .with_options(call_chain_layout.after_attribute())
+                                .fmt(f)?;
                         }
                         _ => {
                             value.format().with_options(Parentheses::Never).fmt(f)?;
@@ -105,8 +112,30 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
             // Allow the `.` on its own line if this is a fluent call chain
             // and the value either requires parenthesizing or is a call or subscript expression
             // (it's a fluent chain but not the first element).
-            else if call_chain_layout == CallChainLayout::Fluent {
+            //
-                if parenthesize_value || value.is_call_expr() || value.is_subscript_expr() {
+            // In preview we also break _at_ the first call in the chain.
+            // For example:
+            //
+            // ```diff
+            // # stable formatting vs. preview
+            //  x = (
+            // -    df.merge()
+            // +    df
+            // +    .merge()
+            //      .groupby()
+            //      .agg()
+            //      .filter()
+            //  )
+            // ```
+            else if matches!(call_chain_layout, CallChainLayout::Fluent(_)) {
+                if parenthesize_value
+                    || value.is_call_expr()
+                    || value.is_subscript_expr()
+                    // Remember to update the doc-comment above when
+                    // stabilizing this behavior.
+                    || (is_fluent_layout_split_first_call_enabled(f.context())
+                        && call_chain_layout.is_first_call_like())
+                {
                     soft_line_break().fmt(f)?;
                 }
             }
@@ -149,7 +178,7 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
         });
 
         let is_call_chain_root = self.call_chain_layout == CallChainLayout::Default
-            && call_chain_layout == CallChainLayout::Fluent;
+            && matches!(call_chain_layout, CallChainLayout::Fluent(_));
         if is_call_chain_root {
             write!(f, [group(&format_inner)])
         } else {
@@ -165,12 +194,15 @@ impl NeedsParentheses for ExprAttribute {
         context: &PyFormatContext,
     ) -> OptionalParentheses {
         // Checks if there are any own line comments in an attribute chain (a.b.c).
-        if CallChainLayout::from_expression(
+        if matches!(
-            self.into(),
+            CallChainLayout::from_expression(
-            context.comments().ranges(),
+                self.into(),
-            context.source(),
+                context.comments().ranges(),
-        ) == CallChainLayout::Fluent
+                context.source(),
-        {
+                context
+            ),
+            CallChainLayout::Fluent(_)
+        ) {
             OptionalParentheses::Multiline
         } else if context.comments().has_dangling(self) {
             OptionalParentheses::Always

crates/ruff_python_formatter/src/expression/expr_call.rs

@@ -47,7 +47,10 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
                 func.format().with_options(Parentheses::Always).fmt(f)
             } else {
                 match func.as_ref() {
-                    Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Attribute(expr) => expr
+                        .format()
+                        .with_options(call_chain_layout.call_like_attribute())
+                        .fmt(f),
                     Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f),
                     Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f),
                     _ => func.format().with_options(Parentheses::Never).fmt(f),
@@ -67,7 +70,7 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
         //     queryset.distinct().order_by(field.name).values_list(field_name_flat_long_long=True)
         // )
         // ```
-        if call_chain_layout == CallChainLayout::Fluent
+        if matches!(call_chain_layout, CallChainLayout::Fluent(_))
             && self.call_chain_layout == CallChainLayout::Default
         {
             group(&fmt_func).fmt(f)
@@ -83,12 +86,15 @@ impl NeedsParentheses for ExprCall {
         _parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if CallChainLayout::from_expression(
+        if matches!(
-            self.into(),
+            CallChainLayout::from_expression(
-            context.comments().ranges(),
+                self.into(),
-            context.source(),
+                context.comments().ranges(),
-        ) == CallChainLayout::Fluent
+                context.source(),
-        {
+                context
+            ),
+            CallChainLayout::Fluent(_)
+        ) {
             OptionalParentheses::Multiline
         } else if context.comments().has_dangling(self) {
             OptionalParentheses::Always

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -51,7 +51,10 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
                 value.format().with_options(Parentheses::Always).fmt(f)
             } else {
                 match value.as_ref() {
-                    Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Attribute(expr) => expr
+                        .format()
+                        .with_options(call_chain_layout.call_like_attribute())
+                        .fmt(f),
                     Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f),
                     Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f),
                     _ => value.format().with_options(Parentheses::Never).fmt(f),
@@ -72,7 +75,7 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
         });
 
         let is_call_chain_root = self.call_chain_layout == CallChainLayout::Default
-            && call_chain_layout == CallChainLayout::Fluent;
+            && matches!(call_chain_layout, CallChainLayout::Fluent(_));
         if is_call_chain_root {
             write!(f, [group(&format_inner)])
         } else {
@@ -88,12 +91,15 @@ impl NeedsParentheses for ExprSubscript {
         context: &PyFormatContext,
     ) -> OptionalParentheses {
         {
-            if CallChainLayout::from_expression(
+            if matches!(
-                self.into(),
+                CallChainLayout::from_expression(
-                context.comments().ranges(),
+                    self.into(),
-                context.source(),
+                    context.comments().ranges(),
-            ) == CallChainLayout::Fluent
+                    context.source(),
-            {
+                    context
+                ),
+                CallChainLayout::Fluent(_)
+            ) {
                 OptionalParentheses::Multiline
             } else if is_expression_parenthesized(
                 self.value.as_ref().into(),

crates/ruff_python_formatter/src/expression/mod.rs

@@ -19,7 +19,9 @@ use crate::expression::parentheses::{
     optional_parentheses, parenthesized,
 };
 use crate::prelude::*;
-use crate::preview::is_hug_parens_with_braces_and_square_brackets_enabled;
+use crate::preview::{
+    is_fluent_layout_more_often_enabled, is_hug_parens_with_braces_and_square_brackets_enabled,
+};
 
 mod binary_like;
 pub(crate) mod expr_attribute;
@@ -883,21 +885,62 @@ pub enum CallChainLayout {
     Default,
 
     /// A nested call chain element that uses fluent style.
-    Fluent,
+    Fluent(AttributeState),
 
     /// A nested call chain element not using fluent style.
     NonFluent,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum AttributeState {
+    CallsOrSubscriptsPreceding(u32),
+    FirstCallOrSubscript,
+    BeforeFirstCallOrSubscript,
+}
+
 impl CallChainLayout {
+    pub(crate) fn call_like_attribute(self) -> Self {
+        match self {
+            Self::Fluent(AttributeState::CallsOrSubscriptsPreceding(x)) => {
+                if x > 1 {
+                    Self::Fluent(AttributeState::CallsOrSubscriptsPreceding(x - 1))
+                } else {
+                    Self::Fluent(AttributeState::FirstCallOrSubscript)
+                }
+            }
+            _ => self,
+        }
+    }
+
+    pub(crate) fn after_attribute(self) -> Self {
+        match self {
+            Self::Fluent(AttributeState::FirstCallOrSubscript) => {
+                Self::Fluent(AttributeState::BeforeFirstCallOrSubscript)
+            }
+            _ => self,
+        }
+    }
+
+    pub(crate) fn is_first_call_like(self) -> bool {
+        matches!(self, Self::Fluent(AttributeState::FirstCallOrSubscript))
+    }
+
     pub(crate) fn from_expression(
         mut expr: ExprRef,
         comment_ranges: &CommentRanges,
         source: &str,
+        // This can be deleted once the preview style
+        // is stabilized
+        context: &PyFormatContext,
     ) -> Self {
+        let mut call_like_count = 0;
+        let mut was_in_call_like = false;
         let mut first_attr_value_parenthesized = false;
 
+        // We can delete this and its uses below once
+        // the preview style [] is stabilized.
         let mut attributes_after_parentheses = 0;
+
         loop {
             match expr {
                 ExprRef::Attribute(ast::ExprAttribute { value, .. }) => {
@@ -907,6 +950,7 @@ impl CallChainLayout {
                     // data[:100].T
                     // ^^^^^^^^^^ value
                     // ```
+                    call_like_count += u32::from(was_in_call_like);
                     if is_expression_parenthesized(value.into(), comment_ranges, source) {
                         // `(a).b`. We preserve these parentheses so don't recurse
                         first_attr_value_parenthesized = true;
@@ -914,7 +958,7 @@ impl CallChainLayout {
                     } else if matches!(value.as_ref(), Expr::Call(_) | Expr::Subscript(_)) {
                         attributes_after_parentheses += 1;
                     }
-
+                    was_in_call_like = false;
                     expr = ExprRef::from(value.as_ref());
                 }
                 // ```
@@ -934,9 +978,16 @@ impl CallChainLayout {
                         break;
                     }
 
+                    was_in_call_like = true;
                     expr = ExprRef::from(inner.as_ref());
                 }
                 _ => {
+                    // We count the first call in the chain even
+                    // if it is not part of an attribute, e.g.
+                    //
+                    // f().g()
+                    // ^^^ count this
+                    call_like_count += u32::from(was_in_call_like);
                     break;
                 }
             }
@@ -945,7 +996,7 @@ impl CallChainLayout {
         if attributes_after_parentheses + u32::from(first_attr_value_parenthesized) < 2 {
             CallChainLayout::NonFluent
         } else {
-            CallChainLayout::Fluent
+            CallChainLayout::Fluent(AttributeState::CallsOrSubscriptsPreceding(call_like_count))
         }
     }
 
@@ -963,12 +1014,13 @@ impl CallChainLayout {
                         item.into(),
                         f.context().comments().ranges(),
                         f.context().source(),
+                        f.context(),
                     )
                 } else {
                     CallChainLayout::NonFluent
                 }
             }
-            layout @ (CallChainLayout::Fluent | CallChainLayout::NonFluent) => layout,
+            layout @ (CallChainLayout::Fluent(_) | CallChainLayout::NonFluent) => layout,
         }
     }
 }

crates/ruff_python_formatter/src/preview.rs

@@ -59,3 +59,17 @@ pub(crate) const fn is_avoid_parens_for_long_as_captures_enabled(
 pub(crate) const fn is_parenthesize_lambda_bodies_enabled(context: &PyFormatContext) -> bool {
     context.is_preview()
 }
+
+/// Returns `true` if the
+/// [`fluent_layout_split_first_call`]() preview
+/// style is enabled.
+pub(crate) const fn is_fluent_layout_split_first_call_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}
+
+/// Returns `true` if the
+/// [`fluent_layout_more_often`]() preview
+/// style is enabled.
+pub(crate) const fn is_fluent_layout_more_often_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}