Call chain formatting in fluent style

Author: konstin

crates/ruff_python_formatter/resources/test/fixtures/ruff/call_chains.py

@@ -0,0 +1,82 @@
+# Test cases for call chains and optional parentheses, with and without fluent style
+
+raise OsError("") from a.aaaaa(
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa
+).a(aaaa)
+
+raise OsError(
+    "sökdjffffsldkfjlhsakfjhalsökafhsöfdahsödfjösaaksjdllllllllllllll"
+) from a.aaaaa(
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa
+).a(
+    aaaa
+)
+
+blogs1 = Blog.objects.filter(entry__headline__contains="Lennon").filter(
+    entry__pub_date__year=2008
+)
+
+blogs2 = Blog.objects.filter(
+    entry__headline__contains="Lennon",
+).filter(
+    entry__pub_date__year=2008,
+)
+
+raise OsError("") from (
+    Blog.objects.filter(
+        entry__headline__contains="Lennon",
+    )
+    .filter(
+        entry__pub_date__year=2008,
+    )
+    .filter(
+        entry__pub_date__year=2008,
+    )
+)
+
+raise OsError("sökdjffffsldkfjlhsakfjhalsökafhsöfdahsödfjösaaksjdllllllllllllll") from (
+    Blog.objects.filter(
+        entry__headline__contains="Lennon",
+    )
+    .filter(
+        entry__pub_date__year=2008,
+    )
+    .filter(
+        entry__pub_date__year=2008,
+    )
+)
+
+# Break only after calls and indexing
+result = (
+    session.query(models.Customer.id)
+    .filter(
+        models.Customer.account_id == account_id, models.Customer.email == email_address
+    )
+    .count()
+)
+
+raise (
+    Blog.objects.filter(
+        entry__headline__contains="Lennon",
+    )
+    .limit_results[:10]
+    .filter(
+        entry__pub_date__month=10,
+    )
+)
+
+# Nested call chains
+blogs = (
+    Blog.objects.filter(
+        entry__headline__contains="Lennon",
+    ).filter(
+        entry__pub_date__year=2008,
+    )
+    + Blog.objects.filter(
+        entry__headline__contains="McCartney",
+    )
+    .limit_results[:10]
+    .filter(
+        entry__pub_date__year=2010,
+    )
+).all()

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/call.py

@@ -52,7 +52,7 @@ def f(*args, **kwargs):
     hey_this_is_a_very_long_call=1, it_has_funny_attributes_asdf_asdf=1, too_long_for_the_line=1, really=True
 )
 
-# TODO(konstin): Call chains/fluent interface (https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#call-chains)
+# Call chains/fluent interface (https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#call-chains)
 result = (
     session.query(models.Customer.id)
     .filter(

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -1,15 +1,26 @@
-use ruff_python_ast::{Constant, Expr, ExprAttribute, ExprConstant};
-
-use ruff_formatter::write;
+use ruff_formatter::{write, FormatRuleWithOptions};
 use ruff_python_ast::node::AnyNodeRef;
+use ruff_python_ast::{Constant, Expr, ExprAttribute, ExprConstant};
 
 use crate::comments::{leading_comments, trailing_comments};
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
 use crate::FormatNodeRule;
 
 #[derive(Default)]
-pub struct FormatExprAttribute;
+pub struct FormatExprAttribute {
+    /// Parentheses for fluent style
+    parentheses: Option<Parentheses>,
+}
+
+impl FormatRuleWithOptions<ExprAttribute, PyFormatContext<'_>> for FormatExprAttribute {
+    type Options = Option<Parentheses>;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
     fn fmt_fields(&self, item: &ExprAttribute, f: &mut PyFormatter) -> FormatResult<()> {
@@ -37,11 +48,18 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
 
         if needs_parentheses {
             value.format().with_options(Parentheses::Always).fmt(f)?;
-        } else if let Expr::Attribute(expr_attribute) = value.as_ref() {
-            // We're in a attribute chain (`a.b.c`). The outermost node adds parentheses if
-            // required, the inner ones don't need them so we skip the `Expr` formatting that
-            // normally adds the parentheses.
-            expr_attribute.format().fmt(f)?;
+        } else if self.parentheses == Some(Parentheses::FluentStyle) {
+            // Fluent style: We need to pass the parentheses on to inner attributes or call chains
+            value
+                .format()
+                .with_options(Parentheses::FluentStyle)
+                .fmt(f)?;
+            match value.as_ref() {
+                Expr::Call(_) | Expr::Subscript(_) => {
+                    soft_line_break().fmt(f)?;
+                }
+                _ => {}
+            }
         } else {
             value.format().fmt(f)?;
         }
@@ -50,16 +68,51 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
             hard_line_break().fmt(f)?;
         }
 
-        write!(
-            f,
-            [
-                text("."),
-                trailing_comments(trailing_dot_comments),
-                (!leading_attribute_comments.is_empty()).then_some(hard_line_break()),
-                leading_comments(leading_attribute_comments),
-                attr.format()
-            ]
-        )
+        if self.parentheses == Some(Parentheses::FluentStyle) {
+            // Fluent style has line breaks before the dot
+            // ```python
+            // blogs3 = (
+            //     Blog.objects.filter(
+            //         entry__headline__contains="Lennon",
+            //     )
+            //     .filter(
+            //         entry__pub_date__year=2008,
+            //     )
+            //     .filter(
+            //         entry__pub_date__year=2008,
+            //     )
+            // )
+            // ```
+            write!(
+                f,
+                [
+                    (!leading_attribute_comments.is_empty()).then_some(hard_line_break()),
+                    leading_comments(leading_attribute_comments),
+                    text("."),
+                    trailing_comments(trailing_dot_comments),
+                    attr.format()
+                ]
+            )
+        } else {
+            // Regular style
+            // ```python
+            // blogs2 = Blog.objects.filter(
+            //     entry__headline__contains="Lennon",
+            // ).filter(
+            //     entry__pub_date__year=2008,
+            // )
+            // ```
+            write!(
+                f,
+                [
+                    text("."),
+                    trailing_comments(trailing_dot_comments),
+                    (!leading_attribute_comments.is_empty()).then_some(hard_line_break()),
+                    leading_comments(leading_attribute_comments),
+                    attr.format()
+                ]
+            )
+        }
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/expression/expr_call.rs

@@ -1,11 +1,10 @@
+use ruff_formatter::{write, FormatRuleWithOptions};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{Expr, ExprCall, Ranged};
+use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::{TextRange, TextSize};
 
 use crate::builders::empty_parenthesized_with_dangling_comments;
-use ruff_formatter::write;
-use ruff_python_ast::node::AnyNodeRef;
-use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
-
 use crate::expression::expr_generator_exp::GeneratorExpParentheses;
 use crate::expression::parentheses::{
     parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
@@ -14,7 +13,19 @@ use crate::prelude::*;
 use crate::FormatNodeRule;
 
 #[derive(Default)]
-pub struct FormatExprCall;
+pub struct FormatExprCall {
+    /// Parentheses for fluent style
+    parentheses: Option<Parentheses>,
+}
+
+impl FormatRuleWithOptions<ExprCall, PyFormatContext<'_>> for FormatExprCall {
+    type Options = Option<Parentheses>;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExprCall> for FormatExprCall {
     fn fmt_fields(&self, item: &ExprCall, f: &mut PyFormatter) -> FormatResult<()> {
@@ -25,28 +36,36 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
             keywords,
         } = item;
 
+        if self.parentheses == Some(Parentheses::FluentStyle) {
+            // Fluent style: We need to pass the parentheses on to inner attributes or call chains
+            func.format()
+                .with_options(Parentheses::FluentStyle)
+                .fmt(f)?;
+        } else {
+            func.format().fmt(f)?;
+        }
+
         // We have a case with `f()` without any argument, which is a special case because we can
         // have a comment with no node attachment inside:
         // ```python
         // f(
         //      # This function has a dangling comment
         // )
         // ```
+        let comments = f.context().comments().clone();
         if args.is_empty() && keywords.is_empty() {
-            let comments = f.context().comments().clone();
             return write!(
                 f,
-                [
-                    func.format(),
-                    empty_parenthesized_with_dangling_comments(
-                        text("("),
-                        comments.dangling_comments(item),
-                        text(")"),
-                    )
-                ]
+                [empty_parenthesized_with_dangling_comments(
+                    text("("),
+                    comments.dangling_comments(item),
+                    text(")"),
+                )]
             );
         }
 
+        debug_assert!(comments.dangling_comments(item).is_empty());
+
         let all_args = format_with(|f: &mut PyFormatter| {
             let source = f.context().source();
             let mut joiner = f.join_comma_separated(item.end());
@@ -88,7 +107,6 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
         write!(
             f,
             [
-                func.format(),
                 // The outer group is for things like
                 // ```python
                 // get_collection(
@@ -104,7 +122,6 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
                 //     hey_this_is_a_very_long_call, it_has_funny_attributes_asdf_asdf, really=True
                 // )
                 // ```
-                // TODO(konstin): Doesn't work see wrongly formatted test
                 parenthesized("(", &group(&all_args), ")")
             ]
         )

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -1,18 +1,29 @@
-use ruff_python_ast::{Expr, ExprSubscript};
-
-use ruff_formatter::{format_args, write};
+use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
+use ruff_python_ast::{Expr, ExprSubscript};
 
 use crate::comments::trailing_comments;
 use crate::context::NodeLevel;
 use crate::context::PyFormatContext;
 use crate::expression::expr_tuple::TupleParentheses;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
 use crate::FormatNodeRule;
 
 #[derive(Default)]
-pub struct FormatExprSubscript;
+pub struct FormatExprSubscript {
+    parentheses: Option<Parentheses>,
+}
+
+impl FormatRuleWithOptions<ExprSubscript, PyFormatContext<'_>> for FormatExprSubscript {
+    /// Parentheses for fluent style
+    type Options = Option<Parentheses>;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
     fn fmt_fields(&self, item: &ExprSubscript, f: &mut PyFormatter) -> FormatResult<()> {
@@ -30,15 +41,24 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
             "The subscript expression must have at most a single comment, the one after the bracket"
         );
 
+        let format_value = format_with(|f| {
+            if self.parentheses == Some(Parentheses::FluentStyle) {
+                // Fluent style: We need to pass the parentheses on to inner attributes or call chains
+                value.format().with_options(Parentheses::FluentStyle).fmt(f)
+            } else {
+                value.format().fmt(f)
+            }
+        });
+
         if let NodeLevel::Expression(Some(group_id)) = f.context().node_level() {
             // Enforce the optional parentheses for parenthesized values.
             f.context_mut().set_node_level(NodeLevel::Expression(None));
-            let result = value.format().fmt(f);
+            let result = format_value.fmt(f);
             f.context_mut()
                 .set_node_level(NodeLevel::Expression(Some(group_id)));
             result?;
         } else {
-            value.format().fmt(f)?;
+            format_value.fmt(f)?;
         }
 
         let format_slice = format_with(|f: &mut PyFormatter| {