Encode in parent

Author: charliermarsh

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(),

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`

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)]

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,15 @@ 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 +45,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 +64,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

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