Implement Black's parenthesize-long-type-hints preview style

astral-sh · Dec 12, 2023 · 104c63d · 104c63d
1 parent ed58d9b
commit 104c63d
Show file tree

Hide file tree

Showing 7 changed files with 113 additions and 199 deletions.
diff --git a/crates/ruff_python_formatter/src/other/parameter.rs b/crates/ruff_python_formatter/src/other/parameter.rs
@@ -2,6 +2,7 @@ use ruff_formatter::write;
 use ruff_python_ast::Parameter;
 
 use crate::prelude::*;
+use crate::statement::stmt_ann_assign::FormatAnnotation;
 
 #[derive(Default)]
 pub struct FormatParameter;
@@ -17,7 +18,7 @@ impl FormatNodeRule<Parameter> for FormatParameter {
         name.format().fmt(f)?;
 
         if let Some(annotation) = annotation {
-            write!(f, [token(":"), space(), annotation.format()])?;
+            write!(f, [token(":"), space(), FormatAnnotation::new(annotation)])?;
         }
 
         Ok(())

diff --git a/crates/ruff_python_formatter/src/preview.rs b/crates/ruff_python_formatter/src/preview.rs
@@ -24,3 +24,9 @@ pub(crate) const fn is_prefer_splitting_right_hand_side_of_assignments_enabled(
 ) -> bool {
     context.is_preview()
 }
+
+/// FIXME ADD proper URl once I have internet again.
+/// Returns `true` if the [`parenthesize_long_type_hints`](TOOD) preview style is enabled.
+pub(crate) const fn is_parenthesize_long_type_hints_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}
diff --git a/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs b/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs
@@ -1,10 +1,15 @@
+use crate::builders::parenthesize_if_expands;
 use ruff_formatter::write;
-use ruff_python_ast::StmtAnnAssign;
+use ruff_python_ast::{Expr, StmtAnnAssign};
 
 use crate::comments::{SourceComment, SuppressionKind};
-use crate::expression::has_parentheses;
+use crate::expression::parentheses::Parentheses;
+use crate::expression::{has_own_parentheses, has_parentheses};
 use crate::prelude::*;
-use crate::preview::is_prefer_splitting_right_hand_side_of_assignments_enabled;
+use crate::preview::{
+    is_parenthesize_long_type_hints_enabled,
+    is_prefer_splitting_right_hand_side_of_assignments_enabled,
+};
 use crate::statement::stmt_assign::{
     AnyAssignmentOperator, AnyBeforeOperator, FormatStatementsLastExpression,
 };
@@ -27,7 +32,9 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
 
         if let Some(value) = value {
             if is_prefer_splitting_right_hand_side_of_assignments_enabled(f.context())
-                && has_parentheses(annotation, f.context()).is_some()
+                && (has_parentheses(annotation, f.context()).is_some()
+                    || (is_parenthesize_long_type_hints_enabled(f.context()))
+                        && should_parenthesize_annotation(annotation, f.context()))
             {
                 FormatStatementsLastExpression::RightToLeft {
                     before_operator: AnyBeforeOperator::Expression(annotation),
@@ -49,7 +56,7 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                 )?;
             }
         } else {
-            annotation.format().fmt(f)?;
+            FormatAnnotation::new(annotation).fmt(f)?;
         }
 
         if f.options().source_type().is_ipynb()
@@ -71,3 +78,45 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
         SuppressionKind::has_skip_comment(trailing_comments, context.source())
     }
 }
+
+pub(crate) struct FormatAnnotation<'a> {
+    annotation: &'a Expr,
+}
+
+impl<'a> FormatAnnotation<'a> {
+    pub(crate) fn new(annotation: &'a Expr) -> Self {
+        Self { annotation }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatAnnotation<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        if is_parenthesize_long_type_hints_enabled(f.context()) {
+            let comments = f.context().comments();
+            if comments.has_leading(self.annotation) || comments.has_trailing(self.annotation) {
+                self.annotation
+                    .format()
+                    .with_options(Parentheses::Always)
+                    .fmt(f)
+            } else if should_parenthesize_annotation(self.annotation, f.context()) {
+                parenthesize_if_expands(&self.annotation.format().with_options(Parentheses::Never))
+                    .fmt(f)
+            } else {
+                self.annotation
+                    .format()
+                    .with_options(Parentheses::Never)
+                    .fmt(f)
+            }
+        } else {
+            self.annotation.format().fmt(f)
+        }
+    }
+}
+
+/// Returns `true` if an annotation should be parenthesized if it splits over multiple lines.
+pub(crate) fn should_parenthesize_annotation(
+    annotation: &Expr,
+    context: &PyFormatContext<'_>,
+) -> bool {
+    !matches!(annotation, Expr::Name(_)) && has_own_parentheses(annotation, context).is_none()
+}
diff --git a/...rmatter/tests/snapshots/black_compatibility@cases__pep604_union_types_line_breaks.py.snap b/...rmatter/tests/snapshots/black_compatibility@cases__pep604_union_types_line_breaks.py.snap
@@ -95,62 +95,7 @@ def f(
 ```diff
 --- Black
 +++ Ruff
-@@ -7,23 +7,13 @@
- )
- 
- # "AnnAssign"s now also work
--z: (
--    Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--)
--z: Short | Short2 | Short3 | Short4
--z: int
--z: int
-+z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong
-+z: (Short | Short2 | Short3 | Short4)
-+z: (int)
-+z: (int)
- 
- 
--z: (
--    Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--    | Loooooooooooooooooooooooong
--) = 7
-+z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong = 7
- z: Short | Short2 | Short3 | Short4 = 8
- z: int = 2.3
- z: int = foo()
-@@ -63,7 +53,7 @@
- 
- 
- # remove unnecessary paren
--def foo(i: int) -> None: ...
-+def foo(i: (int)) -> None: ...
- 
- 
- # this is a syntax error in the type annotation according to mypy, but it's not invalid *python* code, so make sure we don't mess with it and make it so.
-@@ -72,12 +62,10 @@
- 
- def foo(
-     i: int,
--    x: (
--        Loooooooooooooooooooooooong
--        | Looooooooooooooooong
--        | Looooooooooooooooooooong
--        | Looooooong
--    ),
-+    x: Loooooooooooooooooooooooong
-+    | Looooooooooooooooong
-+    | Looooooooooooooooooooong
-+    | Looooooong,
-     *,
-     s: str,
- ) -> None:
-@@ -88,7 +76,7 @@
+@@ -88,7 +88,7 @@
  async def foo(
      q: str | None = Query(
          None, title="Some long title", description="Some long description"
@@ -173,13 +118,23 @@ z = (
 )
 
 # "AnnAssign"s now also work
-z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong
-z: (Short | Short2 | Short3 | Short4)
-z: (int)
-z: (int)
+z: (
+    Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+)
+z: Short | Short2 | Short3 | Short4
+z: int
+z: int
 
 
-z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong = 7
+z: (
+    Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+) = 7
 z: Short | Short2 | Short3 | Short4 = 8
 z: int = 2.3
 z: int = foo()
@@ -219,7 +174,7 @@ x = (
 
 
 # remove unnecessary paren
-def foo(i: (int)) -> None: ...
+def foo(i: int) -> None: ...
 
 
 # this is a syntax error in the type annotation according to mypy, but it's not invalid *python* code, so make sure we don't mess with it and make it so.
@@ -228,10 +183,12 @@ def foo(i: (int,)) -> None: ...
 
 def foo(
     i: int,
-    x: Loooooooooooooooooooooooong
-    | Looooooooooooooooong
-    | Looooooooooooooooooooong
-    | Looooooong,
+    x: (
+        Loooooooooooooooooooooooong
+        | Looooooooooooooooong
+        | Looooooooooooooooooooong
+        | Looooooong
+    ),
     *,
     s: str,
 ) -> None:

diff --git a/...tests/snapshots/black_compatibility@cases__preview_long_strings__type_annotations.py.snap b/...tests/snapshots/black_compatibility@cases__preview_long_strings__type_annotations.py.snap
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__ann_assign.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__ann_assign.py.snap
@@ -84,6 +84,16 @@ class DefaultRunner:
  
  JSONSerializable: TypeAlias = (
      "str | int | float | bool | None | list | tuple | JSONMapping"
+@@ -29,6 +29,6 @@
+ 
+ # Regression test: Don't forget the parentheses in the annotation when breaking
+ class DefaultRunner:
+-    task_runner_cls: TaskRunnerProtocol | typing.Callable[
+-        [], typing.Any
+-    ] = DefaultTaskRunner
++    task_runner_cls: TaskRunnerProtocol | typing.Callable[[], typing.Any] = (
++        DefaultTaskRunner
++    )
 ```