Parenthesize long type annotations in annotated assignment statements

Author: MichaReiser

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/long_type_annotations.py

@@ -0,0 +1,157 @@
+x1: A[b] | EventHandler | EventSpec | list[EventHandler | EventSpec] | Other | More | AndMore | None = None
+
+x2: "VeryLongClassNameWithAwkwardGenericSubtype[int] |" "VeryLongClassNameWithAwkwardGenericSubtype[str]"
+
+x6: VeryLongClassNameWithAwkwardGenericSubtype[
+    integeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeer,
+    VeryLongClassNameWithAwkwardGenericSubtype,
+    str
+] = True
+
+
+x7: CustomTrainingJob | CustomContainerTrainingJob | CustomPythonPackageTrainingJob
+x8: (
+    None
+    | datasets.ImageDataset
+    | datasets.TabularDataset
+    | datasets.TextDataset
+    | datasets.VideoDataset
+) = None
+
+x9: None | (
+    datasets.ImageDataset
+    | datasets.TabularDataset
+    | datasets.TextDataset
+    | datasets.VideoDataset
+) = None
+
+
+x10: (
+    aaaaaaaaaaaaaaaaaaaaaaaa[
+        bbbbbbbbbbb,
+        Subscript
+        | None
+        | datasets.ImageDataset
+        | datasets.TabularDataset
+        | datasets.TextDataset
+        | datasets.VideoDataset,
+    ],
+    bbb[other],
+) = None
+
+x11: None | [
+    datasets.ImageDataset,
+    datasets.TabularDataset,
+    datasets.TextDataset,
+    datasets.VideoDataset,
+] = None
+
+x12: None | [
+        datasets.ImageDataset,
+        datasets.TabularDataset,
+        datasets.TextDataset,
+        datasets.VideoDataset,
+    ] | Other = None
+
+
+x13: [
+     datasets.ImageDataset,
+     datasets.TabularDataset,
+     datasets.TextDataset,
+     datasets.VideoDataset,
+] | Other = None
+
+x14: [
+    datasets.ImageDataset,
+    datasets.TabularDataset,
+    datasets.TextDataset,
+    datasets.VideoDataset,
+] | [
+     datasets.ImageDataset,
+     datasets.TabularDataset,
+     datasets.TextDataset,
+     datasets.VideoDataset,
+] = None
+
+x15: [
+    datasets.ImageDataset,
+    datasets.TabularDataset,
+    datasets.TextDataset,
+    datasets.VideoDataset,
+] | [
+    datasets.ImageDataset,
+    datasets.TabularDataset,
+    datasets.TextDataset,
+    datasets.VideoDataset,
+] | Other = None
+
+x16: None | Literal[
+    "split",
+    "a bit longer",
+    "records",
+    "index",
+    "table",
+    "columns",
+    "values",
+] = None
+
+x17: None | [
+    datasets.ImageDataset,
+    datasets.TabularDataset,
+    datasets.TextDataset,
+    datasets.VideoDataset,
+]
+
+
+class Test:
+    safe_age: Decimal  #  the user's age, used to determine if it's safe for them to use ruff
+    applied_fixes: int  # the number of fixes that this user applied. Used for ranking the users with the most applied fixes.
+    string_annotation: "Test"  # a long comment after a quoted, runtime-only type annotation
+
+
+##########
+# Comments
+
+leading: (
+    #  Leading comment
+    None | dataset.ImageDataset
+)
+
+leading_with_value: (
+    #  Leading comment
+    None
+    | dataset.ImageDataset
+) = None
+
+leading_open_parentheses: ( #  Leading comment
+    None
+    | dataset.ImageDataset
+)
+
+leading_open_parentheses_with_value: ( #  Leading comment
+    None
+    | dataset.ImageDataset
+) = None
+
+trailing: (
+    None | dataset.ImageDataset  # trailing comment
+)
+
+trailing_with_value: (
+    None | dataset.ImageDataset  # trailing comment
+) = None
+
+trailing_own_line: (
+    None | dataset.ImageDataset
+    #  trailing own line
+)
+
+trailing_own_line_with_value: (
+    None | dataset.ImageDataset
+    #  trailing own line
+) = None
+
+nested_comment: None | [
+    # a list of strings
+    str
+] = None

crates/ruff_python_formatter/src/expression/mod.rs

@@ -529,7 +529,7 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for Expr {
 ///
 /// This mimics Black's [`_maybe_split_omitting_optional_parens`](https://github.com/psf/black/blob/d1248ca9beaf0ba526d265f4108836d89cf551b7/src/black/linegen.py#L746-L820)
 #[allow(clippy::if_same_then_else)]
-fn can_omit_optional_parentheses(expr: &Expr, context: &PyFormatContext) -> bool {
+pub(crate) fn can_omit_optional_parentheses(expr: &Expr, context: &PyFormatContext) -> bool {
     let mut visitor = CanOmitOptionalParenthesesVisitor::new(context);
     visitor.visit_subexpression(expr);
 
@@ -1195,3 +1195,75 @@ impl From<Operator> for OperatorPrecedence {
         }
     }
 }
+
+/// Returns `true` if `expr` is an expression that can be split into multiple lines.
+///
+/// Returns `false` for expressions that are guaranteed to never split.
+pub(crate) fn is_splittable_expression(expr: &Expr, context: &PyFormatContext) -> bool {
+    match expr {
+        // Single token expressions. They never have any split points.
+        Expr::NamedExpr(_)
+        | Expr::Name(_)
+        | Expr::NumberLiteral(_)
+        | Expr::BooleanLiteral(_)
+        | Expr::NoneLiteral(_)
+        | Expr::EllipsisLiteral(_)
+        | Expr::Slice(_)
+        | Expr::IpyEscapeCommand(_) => false,
+
+        // Expressions that insert split points when parenthesized.
+        Expr::Compare(_)
+        | Expr::BinOp(_)
+        | Expr::BoolOp(_)
+        | Expr::IfExp(_)
+        | Expr::GeneratorExp(_)
+        | Expr::Subscript(_)
+        | Expr::Await(_)
+        | Expr::ListComp(_)
+        | Expr::SetComp(_)
+        | Expr::DictComp(_)
+        | Expr::YieldFrom(_) => true,
+
+        // Sequence types can split if they contain at least one element.
+        Expr::Tuple(tuple) => !tuple.elts.is_empty(),
+        Expr::Dict(dict) => !dict.values.is_empty(),
+        Expr::Set(set) => !set.elts.is_empty(),
+        Expr::List(list) => !list.elts.is_empty(),
+
+        Expr::UnaryOp(unary) => is_splittable_expression(unary.operand.as_ref(), context),
+        Expr::Yield(ast::ExprYield { value, .. }) => value.is_some(),
+
+        Expr::Call(ast::ExprCall {
+            arguments, func, ..
+        }) => {
+            !arguments.is_empty()
+                || is_expression_parenthesized(
+                    func.as_ref().into(),
+                    context.comments().ranges(),
+                    context.source(),
+                )
+        }
+
+        // String like literals can expand if they are implicit concatenated.
+        Expr::FString(fstring) => fstring.value.is_implicit_concatenated(),
+        Expr::StringLiteral(string) => string.value.is_implicit_concatenated(),
+        Expr::BytesLiteral(bytes) => bytes.value.is_implicit_concatenated(),
+
+        // Expressions that have no split points per se, but they contain nested sub expressions that might expand.
+        Expr::Lambda(ast::ExprLambda {
+            body: expression, ..
+        })
+        | Expr::Starred(ast::ExprStarred {
+            value: expression, ..
+        })
+        | Expr::Attribute(ast::ExprAttribute {
+            value: expression, ..
+        }) => {
+            is_expression_parenthesized(
+                expression.into(),
+                context.comments().ranges(),
+                context.source(),
+            ) || is_splittable_expression(expression.as_ref(), context)
+        }
+    }
+}

crates/ruff_python_formatter/src/preview.rs

@@ -25,6 +25,11 @@ pub(crate) const fn is_prefer_splitting_right_hand_side_of_assignments_enabled(
     context.is_preview()
 }
 
+/// Returns `true` if the [`parenthesize_long_type_hints`](https://github.com/astral-sh/ruff/issues/8894) preview style is enabled.
+pub(crate) const fn is_parenthesize_long_type_hints_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}
+
 /// Returns `true` if the [`no_blank_line_before_class_docstring`] preview style is enabled.
 ///
 /// [`no_blank_line_before_class_docstring`]: https://github.com/astral-sh/ruff/issues/8888

crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs

@@ -2,9 +2,13 @@ use ruff_formatter::write;
 use ruff_python_ast::StmtAnnAssign;
 
 use crate::comments::{SourceComment, SuppressionKind};
-use crate::expression::has_parentheses;
+use crate::expression::parentheses::Parentheses;
+use crate::expression::{has_parentheses, is_splittable_expression};
 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 +31,11 @@ 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()
+                // The `has_parentheses` check can be removed when stabilizing `is_parenthesize_long_type_hints`.
+                // because `is_splittable_expression` covers both.
+                && (has_parentheses(annotation, f.context()).is_some()
+                    || (is_parenthesize_long_type_hints_enabled(f.context())
+                        && is_splittable_expression(annotation, f.context())))
             {
                 FormatStatementsLastExpression::RightToLeft {
                     before_operator: AnyBeforeOperator::Expression(annotation),
@@ -37,10 +45,28 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                 }
                 .fmt(f)?;
             } else {
+                // Remove unnecessary parentheses around the annotation if the parenthesize long type hints preview style is enabled.
+                // Ensure we keep the parentheses if the annotation has any comments.
+                if is_parenthesize_long_type_hints_enabled(f.context()) {
+                    if f.context().comments().has_leading(annotation.as_ref())
+                        || f.context().comments().has_trailing(annotation.as_ref())
+                    {
+                        annotation
+                            .format()
+                            .with_options(Parentheses::Always)
+                            .fmt(f)?;
+                    } else {
+                        annotation
+                            .format()
+                            .with_options(Parentheses::Never)
+                            .fmt(f)?;
+                    }
+                } else {
+                    annotation.format().fmt(f)?;
+                }
                 write!(
                     f,
                     [
-                        annotation.format(),
                         space(),
                         token("="),
                         space(),
@@ -49,7 +75,19 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                 )?;
             }
         } else {
-            annotation.format().fmt(f)?;
+            // Parenthesize the value and inline the comment if it is a "simple" type annotation, similar
+            // to what we do with the value.
+            // ```python
+            // class Test:
+            //     safe_age: (
+            //         Decimal  #  the user's age, used to determine if it's safe for them to use ruff
+            //     )
+            // ```
+            if is_parenthesize_long_type_hints_enabled(f.context()) {
+                FormatStatementsLastExpression::left_to_right(annotation, item).fmt(f)?;
+            } else {
+                annotation.format().fmt(f)?;
+            }
         }
 
         if f.options().source_type().is_ipynb()

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -9,11 +9,18 @@ use crate::comments::{
 };
 use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::parentheses::{
-    is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parentheses, Parenthesize,
+    is_expression_parenthesized, optional_parentheses, NeedsParentheses, OptionalParentheses,
+    Parentheses, Parenthesize,
+};
+use crate::expression::{
+    can_omit_optional_parentheses, has_own_parentheses, has_parentheses,
+    maybe_parenthesize_expression,
 };
-use crate::expression::{has_own_parentheses, has_parentheses, maybe_parenthesize_expression};
 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::trailing_semicolon;
 
 #[derive(Default)]
@@ -686,8 +693,17 @@ impl Format<PyFormatContext<'_>> for AnyBeforeOperator<'_> {
                 }
                 // Never parenthesize targets that come with their own parentheses, e.g. don't parenthesize lists or dictionary literals.
                 else if should_parenthesize_target(expression, f.context()) {
-                    parenthesize_if_expands(&expression.format().with_options(Parentheses::Never))
+                    if is_parenthesize_long_type_hints_enabled(f.context())
+                        && can_omit_optional_parentheses(expression, f.context())
+                    {
+                        optional_parentheses(&expression.format().with_options(Parentheses::Never))
+                            .fmt(f)
+                    } else {
+                        parenthesize_if_expands(
+                            &expression.format().with_options(Parentheses::Never),
+                        )
                         .fmt(f)
+                    }
                 } else {
                     expression.format().with_options(Parentheses::Never).fmt(f)
                 }