Skip to content

Preserve type-expression parentheses in the formatter#23867

Merged
charliermarsh merged 3 commits intomainfrom
codex/preserve-type-expression-parens
Mar 11, 2026
Merged

Preserve type-expression parentheses in the formatter#23867
charliermarsh merged 3 commits intomainfrom
codex/preserve-type-expression-parens

Conversation

@charliermarsh
Copy link
Copy Markdown
Member

@charliermarsh charliermarsh commented Mar 10, 2026

Summary

This is a follow-up to #23865 generated by prompted Codex as follows:

Can you find another bug like #23864 in the formatter for us to fix?

Here are some examples that previously had their parentheses stripped (producing syntax errors) but now format correctly:

# Annotated assignments with named expressions
x: (value := int) = 1     # was formatted as: x: value := int = 1

# Annotated assignments with yield
def f():
    y: (yield 1) = 1       # was formatted as: y: yield 1 = 1
    z: (yield 1)            # was formatted as: z: yield 1

# Annotated assignments with await
async def g():
    a: (await g()) = 1      # was formatted as: a: await g() = 1
    b: (await g())           # was formatted as: b: await g()

# Return annotations with await
async def h() -> (await g()):  # was formatted as: async def h() -> await g():
    pass

# Type aliases
type X = (value := int)    # was formatted as: type X = value := int

In each case, the parentheses are required for the expression to be syntactically valid, but the formatter was previously stripping them.

@charliermarsh charliermarsh marked this pull request as draft March 10, 2026 14:38
@charliermarsh charliermarsh force-pushed the codex/preserve-type-expression-parens branch from e50d48b to 893d0f0 Compare March 10, 2026 14:41
@charliermarsh charliermarsh added bug Something isn't working formatter Related to the formatter labels Mar 10, 2026
@astral-sh-bot
Copy link
Copy Markdown

astral-sh-bot bot commented Mar 10, 2026

ruff-ecosystem results

Formatter (stable)

✅ ecosystem check detected no format changes.

Formatter (preview)

✅ ecosystem check detected no format changes.

@charliermarsh charliermarsh force-pushed the codex/preserve-type-expression-parens branch from c254533 to 76bb462 Compare March 10, 2026 15:22
@charliermarsh charliermarsh marked this pull request as ready for review March 10, 2026 16:05
Copy link
Copy Markdown
Member

@MichaReiser MichaReiser left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clever to use codex to find more issues related to type annotations that are syntatically correct but uncommon.

It would be nice if this wouldn't require more special casing in many places and could instead just be handled by the existing needs_parantheses checks, but I don't see a clean way of doing this without increasing the complexity in some of the formatting paths where we don't call needs_parantheses today. So I think what you have here is fine, unless you see a clear way on how we could use needs_parantheses.

simple: _,
} = item;
let comments = f.context().comments().clone();
let preserve_annotation_parentheses = is_invalid_type_expression(annotation);
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to move the check whether the parentheses are needed into the NeedsParentheses implementation of Named, Await, Yield and YieldFrom and inverse the branches in the if (move the needs_parantheses == Always to the top)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you prefer this?

diff --git a/crates/ruff_python_formatter/src/expression/expr_await.rs b/crates/ruff_python_formatter/src/expression/expr_await.rs
index cd862f2f9a..222f5c0a80 100644
--- a/crates/ruff_python_formatter/src/expression/expr_await.rs
+++ b/crates/ruff_python_formatter/src/expression/expr_await.rs
@@ -1,10 +1,12 @@
 use ruff_formatter::write;
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::ExprAwait;
+use ruff_text_size::Ranged;
 
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{
     NeedsParentheses, OptionalParentheses, Parenthesize, is_expression_parenthesized,
+    is_type_annotation_of,
 };
 use crate::prelude::*;
 
@@ -36,7 +38,7 @@ impl NeedsParentheses for ExprAwait {
         parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if parent.is_expr_await() {
+        if parent.is_expr_await() || is_type_annotation_of(self.range(), parent) {
             OptionalParentheses::Always
         } else if is_expression_parenthesized(
             self.value.as_ref().into(),
diff --git a/crates/ruff_python_formatter/src/expression/expr_yield.rs b/crates/ruff_python_formatter/src/expression/expr_yield.rs
index 4b335e208e..8bce330f70 100644
--- a/crates/ruff_python_formatter/src/expression/expr_yield.rs
+++ b/crates/ruff_python_formatter/src/expression/expr_yield.rs
@@ -6,6 +6,7 @@ use ruff_text_size::{Ranged, TextRange};
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{
     NeedsParentheses, OptionalParentheses, Parenthesize, is_expression_parenthesized,
+    is_type_annotation_of,
 };
 use crate::prelude::*;
 
@@ -42,6 +43,10 @@ impl NeedsParentheses for AnyExpressionYield<'_> {
         parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
+        if is_type_annotation_of(self.range(), parent) {
+            return OptionalParentheses::Always;
+        }
+
         // According to https://docs.python.org/3/reference/grammar.html There are two situations
         // where we do not want to always parenthesize a yield expression:
         //  1. Right hand side of an assignment, e.g. `x = yield y`
diff --git a/crates/ruff_python_formatter/src/expression/parentheses.rs b/crates/ruff_python_formatter/src/expression/parentheses.rs
index 18b9b274c1..316a05fc75 100644
--- a/crates/ruff_python_formatter/src/expression/parentheses.rs
+++ b/crates/ruff_python_formatter/src/expression/parentheses.rs
@@ -6,7 +6,7 @@ use ruff_python_trivia::CommentRanges;
 use ruff_python_trivia::{
     BackwardsTokenizer, SimpleToken, SimpleTokenKind, first_non_trivia_token,
 };
-use ruff_text_size::Ranged;
+use ruff_text_size::{Ranged, TextRange};
 
 use crate::comments::{
     SourceComment, dangling_comments, dangling_open_parenthesis_comments, trailing_comments,
@@ -42,6 +42,19 @@ pub(crate) trait NeedsParentheses {
     ) -> OptionalParentheses;
 }
 
+/// Returns `true` if `expr_range` identifies a type annotation child of `parent`,
+/// i.e. the annotation of a [`StmtAnnAssign`] or the return annotation of a [`StmtFunctionDef`].
+pub(crate) fn is_type_annotation_of(expr_range: TextRange, parent: AnyNodeRef) -> bool {
+    match parent {
+        AnyNodeRef::StmtAnnAssign(stmt) => stmt.annotation.range() == expr_range,
+        AnyNodeRef::StmtFunctionDef(stmt) => stmt
+            .returns
+            .as_deref()
+            .is_some_and(|r| r.range() == expr_range),
+        _ => false,
+    }
+}
+
 /// From the perspective of the parent statement or expression, when should the child expression
 /// get parentheses?
 #[derive(Copy, Clone, Debug, PartialEq)]
diff --git a/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs b/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs
index eef5948483..32c6811091 100644
--- a/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs
+++ b/crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs
@@ -1,8 +1,8 @@
 use ruff_formatter::write;
 use ruff_python_ast::StmtAnnAssign;
 
+use crate::expression::is_splittable_expression;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
-use crate::expression::{is_invalid_type_expression, is_splittable_expression};
 use crate::prelude::*;
 use crate::statement::stmt_assign::{
     AnyAssignmentOperator, AnyBeforeOperator, FormatStatementsLastExpression,
@@ -23,12 +23,16 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
             simple: _,
         } = item;
         let comments = f.context().comments().clone();
-        let preserve_annotation_parentheses = is_invalid_type_expression(annotation);
+        let annotation_parentheses =
+            annotation
+                .as_ref()
+                .needs_parentheses(item.into(), f.context());
 
         write!(f, [target.format(), token(":"), space()])?;
 
         if let Some(value) = value {
-            if !preserve_annotation_parentheses && is_splittable_expression(annotation, f.context())
+            if annotation_parentheses != OptionalParentheses::Always
+                && is_splittable_expression(annotation, f.context())
             {
                 FormatStatementsLastExpression::RightToLeft {
                     before_operator: AnyBeforeOperator::Expression(annotation),
@@ -42,15 +46,9 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                 // Ensure we keep the parentheses if the annotation has any comments.
                 let parentheses = if comments.has_leading(annotation.as_ref())
                     || comments.has_trailing(annotation.as_ref())
-                    || matches!(
-                        annotation
-                            .as_ref()
-                            .needs_parentheses(item.into(), f.context()),
-                        OptionalParentheses::Always
-                    ) {
+                    || annotation_parentheses == OptionalParentheses::Always
+                {
                     Parentheses::Always
-                } else if preserve_annotation_parentheses {
-                    Parentheses::Preserve
                 } else {
                     Parentheses::Never
                 };
@@ -67,10 +65,10 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                     ]
                 )?;
             }
-        } else if preserve_annotation_parentheses {
+        } else if annotation_parentheses == OptionalParentheses::Always {
             annotation
                 .format()
-                .with_options(Parentheses::Preserve)
+                .with_options(Parentheses::Always)
                 .fmt(f)?;
         } else {
             // Parenthesize the value and inline the comment if it is a "simple" type annotation, similar
diff --git a/crates/ruff_python_formatter/src/statement/stmt_function_def.rs b/crates/ruff_python_formatter/src/statement/stmt_function_def.rs
index facc9db3c1..61b333c72d 100644
--- a/crates/ruff_python_formatter/src/statement/stmt_function_def.rs
+++ b/crates/ruff_python_formatter/src/statement/stmt_function_def.rs
@@ -1,8 +1,8 @@
 use crate::comments::format::{
     empty_lines_after_leading_comments, empty_lines_before_trailing_comments,
 };
+use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{Parentheses, Parenthesize};
-use crate::expression::{is_invalid_type_expression, maybe_parenthesize_expression};
 use crate::prelude::*;
 use crate::statement::clause::{ClauseHeader, clause};
 use crate::statement::stmt_class_def::FormatDecorators;
@@ -162,11 +162,6 @@ fn format_function_header(f: &mut PyFormatter, item: &StmtFunctionDef) -> Format
                     .format()
                     .with_options(Parentheses::Always)
                     .fmt(f)
-            } else if is_invalid_type_expression(return_annotation) {
-                return_annotation
-                    .format()
-                    .with_options(Parentheses::Preserve)
-                    .fmt(f)
             } else {
                 let parenthesize = if parameters.is_empty() && !comments.has(parameters.as_ref()) {
                     // If the parameters are empty, add parentheses around literal expressions

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From scrolling through the diff, yes. That roughly captures what I had in mind.

@charliermarsh charliermarsh force-pushed the codex/preserve-type-expression-parens branch from 76bb462 to 2b6e1f5 Compare March 11, 2026 18:47
@charliermarsh charliermarsh force-pushed the codex/preserve-type-expression-parens branch from 2b6e1f5 to f4557bd Compare March 11, 2026 19:03
@charliermarsh charliermarsh enabled auto-merge (squash) March 11, 2026 19:03
@charliermarsh charliermarsh merged commit 02f81b1 into main Mar 11, 2026
42 checks passed
@charliermarsh charliermarsh deleted the codex/preserve-type-expression-parens branch March 11, 2026 19:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working formatter Related to the formatter

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants