fix: Use BestFit layout for subscript (#7409)

astral-sh · Sep 16, 2023 · 916dd5b · 916dd5b
1 parent 2cbe173
commit 916dd5b
Show file tree

Hide file tree

Showing 9 changed files with 116 additions and 23 deletions.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/call.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/call.py
@@ -270,3 +270,6 @@ def f(*args, **kwargs):
     f(111111111111111111111111111111111111111111111111111111111111111111111111111111111)
     + 1
 )()
+
+
+result = (object[complicate_caller])("argument").a["b"].test(argument)
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
@@ -54,3 +54,16 @@
         c,
     ]
 ) = 1
+
+def main() -> None:
+    if True:
+        some_very_long_variable_name_abcdefghijk = some_very_long_variable_name_abcdefghijk[
+            some_very_long_variable_name_abcdefghijk.some_very_long_attribute_name
+            == "This is a very long string abcdefghijk"
+        ]
+
+        organization_application = (
+            organization_service.get_organization_applications_by_name(
+                db_request.POST["name"]
+            )
+        )[0]
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/return_annotation.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/return_annotation.py
@@ -180,3 +180,9 @@ def double(a: int) -> (
     int | list[int, int, int, int, int, int, int, int, int, int, int, int, int, int, int, int]  # Hello
 ):
     pass
+
+
+def process_board_action(
+    payload: WildValue, action_type: Optional[str]
+) -> Optional[Tuple[str, str]]:
+    pass
diff --git a/crates/ruff_python_formatter/src/expression/expr_call.rs b/crates/ruff_python_formatter/src/expression/expr_call.rs
@@ -4,7 +4,7 @@ use ruff_python_ast::{Expr, ExprCall};
 
 use crate::comments::{dangling_comments, SourceComment};
 use crate::expression::parentheses::{
-    is_expression_parenthesized, NeedsParentheses, OptionalParentheses,
+    is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::expression::CallChainLayout;
 use crate::prelude::*;
@@ -36,13 +36,17 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
 
         let call_chain_layout = self.call_chain_layout.apply_in_node(item, f);
 
-        let fmt_func = format_with(|f| {
+        let fmt_func = format_with(|f: &mut PyFormatter| {
             // Format the function expression.
-            match func.as_ref() {
-                Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).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().fmt(f),
+            if is_expression_parenthesized(func.into(), f.context().source()) {
+                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::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),
+                }
             }?;
 
             // Format comments between the function and its arguments.

diff --git a/crates/ruff_python_formatter/src/expression/expr_subscript.rs b/crates/ruff_python_formatter/src/expression/expr_subscript.rs
@@ -5,7 +5,7 @@ use ruff_python_ast::{Expr, ExprSubscript};
 use crate::comments::SourceComment;
 use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{
-    is_expression_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses,
+    is_expression_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::expression::CallChainLayout;
 use crate::prelude::*;
@@ -42,13 +42,17 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
             "A subscript expression can only have a single dangling comment, the one after the bracket"
         );
 
-        let format_inner = format_with(|f| {
-            match value.as_ref() {
-                Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).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().fmt(f)?,
-            }
+        let format_inner = format_with(|f: &mut PyFormatter| {
+            if is_expression_parenthesized(value.into(), f.context().source()) {
+                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::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),
+                }
+            }?;
 
             let format_slice = format_with(|f: &mut PyFormatter| {
                 if let Expr::Tuple(tuple) = slice.as_ref() {
@@ -85,7 +89,7 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
 impl NeedsParentheses for ExprSubscript {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
         {
@@ -104,7 +108,21 @@ impl NeedsParentheses for ExprSubscript {
                 OptionalParentheses::Never
             } else {
                 match self.value.needs_parentheses(self.into(), context) {
-                    OptionalParentheses::BestFit => OptionalParentheses::Never,
+                    OptionalParentheses::BestFit => {
+                        if parent.as_stmt_function_def().is_some_and(|function_def| {
+                            function_def
+                                .returns
+                                .as_deref()
+                                .and_then(Expr::as_subscript_expr)
+                                == Some(self)
+                        }) {
+                            // Don't use the best fitting layout for return type annotation because it results in the
+                            // return type expanding before the parameters.
+                            OptionalParentheses::Never
+                        } else {
+                            OptionalParentheses::BestFit
+                        }
+                    }
                     parentheses => parentheses,
                 }
             }

diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs
@@ -217,12 +217,14 @@ if True:
     #[test]
     fn quick_test() {
         let src = r#"
-(header.timecnt * 5  # Transition times and types
-  + header.typecnt * 6  # Local time type records
-  + header.charcnt  # Time zone designations
-  + header.leapcnt * 8  # Leap second records
-  + header.isstdcnt  # Standard/wall indicators
-  + header.isutcnt)  # UT/local indicators
+def main() -> None:
+    if True:
+        some_very_long_variable_name_abcdefghijk = Foo()
+        some_very_long_variable_name_abcdefghijk = some_very_long_variable_name_abcdefghijk[
+            some_very_long_variable_name_abcdefghijk.some_very_long_attribute_name
+            == "This is a very long string abcdefghijk"
+        ]
+
 "#;
         // Tokenize once
         let mut tokens = Vec::new();

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap
@@ -276,6 +276,9 @@ result = (
     f(111111111111111111111111111111111111111111111111111111111111111111111111111111111)
     + 1
 )()
+
+
+result = (object[complicate_caller])("argument").a["b"].test(argument)
 ```
 
 ## Output
@@ -548,6 +551,9 @@ result = (
     f(111111111111111111111111111111111111111111111111111111111111111111111111111111111)
     + 1
 )()
+
+
+result = (object[complicate_caller])("argument").a["b"].test(argument)
 ```
 
 

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
@@ -60,6 +60,19 @@ x = (
         c,
     ]
 ) = 1
+
+def main() -> None:
+    if True:
+        some_very_long_variable_name_abcdefghijk = some_very_long_variable_name_abcdefghijk[
+            some_very_long_variable_name_abcdefghijk.some_very_long_attribute_name
+            == "This is a very long string abcdefghijk"
+        ]
+
+        organization_application = (
+            organization_service.get_organization_applications_by_name(
+                db_request.POST["name"]
+            )
+        )[0]
 ```
 
 ## Output
@@ -122,6 +135,22 @@ x = [  # comment
     b,
     c,
 ] = 1
+
+
+def main() -> None:
+    if True:
+        some_very_long_variable_name_abcdefghijk = (
+            some_very_long_variable_name_abcdefghijk[
+                some_very_long_variable_name_abcdefghijk.some_very_long_attribute_name
+                == "This is a very long string abcdefghijk"
+            ]
+        )
+
+        organization_application = (
+            organization_service.get_organization_applications_by_name(
+                db_request.POST["name"]
+            )
+        )[0]
 ```
 
 

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__return_annotation.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__return_annotation.py.snap
@@ -186,6 +186,12 @@ def double(a: int) -> (
     int | list[int, int, int, int, int, int, int, int, int, int, int, int, int, int, int, int]  # Hello
 ):
     pass
+
+
+def process_board_action(
+    payload: WildValue, action_type: Optional[str]
+) -> Optional[Tuple[str, str]]:
+    pass
 ```
 
 ## Output
@@ -515,6 +521,12 @@ def double(
     ]  # Hello
 ):
     pass
+
+
+def process_board_action(
+    payload: WildValue, action_type: Optional[str]
+) -> Optional[Tuple[str, str]]:
+    pass
 ```
-Original file line number
+Diff line change
@@ Expand Up / @@ -270,3 +270,6 @@ def f(*args, **kwargs): @@
         f(111111111111111111111111111111111111111111111111111111111111111111111111111111111)
         + 1
     )()
+    result = (object[complicate_caller])("argument").a["b"].test(argument)