[pyupgrade] Preserve grouping in multiline UP040 fixes

crates/ruff_linter/resources/test/fixtures/pyupgrade/UP040.py

@@ -132,3 +132,20 @@ class Foo:
 # Test case for TypeVar with default - should be converted when preview mode is enabled
 T_default = TypeVar("T_default", default=int)
 DefaultList: TypeAlias = list[T_default]
+
+# Preserve parentheses around a multiline union in a TypeAliasType call.
+LongAlias = TypeAliasType(
+    "LongAlias",
+    int
+    | str
+    | float
+)
+
+# An already parenthesized value keeps its own parentheses.
+ParenAlias = TypeAliasType(
+    "ParenAlias",
+    (
+        int
+        | str
+    ),
+)

crates/ruff_linter/src/rules/pyupgrade/rules/pep695/non_pep695_type_alias.rs

@@ -5,6 +5,7 @@ use ruff_python_ast::name::Name;
 use ruff_python_ast::token::parenthesized_range;
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{Expr, ExprCall, ExprName, Keyword, StmtAnnAssign, StmtAssign, StmtRef};
+use ruff_python_edits::unwrapped_call_argument;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -136,9 +137,11 @@ pub(crate) fn non_pep695_type_alias_type(checker: &Checker, stmt: &StmtAssign) {
 
     let StmtAssign { targets, value, .. } = stmt;
 
-    let Expr::Call(ExprCall {
-        func, arguments, ..
-    }) = value.as_ref()
+    let Expr::Call(
+        call @ ExprCall {
+            func, arguments, ..
+        },
+    ) = value.as_ref()
     else {
         return;
     };
@@ -195,9 +198,15 @@ pub(crate) fn non_pep695_type_alias_type(checker: &Checker, stmt: &StmtAssign) {
         checker,
         stmt.into(),
         &target_name.id,
-        value,
         &vars,
         TypeAliasKind::TypeAliasType,
+        &unwrapped_call_argument(
+            call,
+            value,
+            Some(stmt.into()),
+            checker.tokens(),
+            checker.source(),
+        ),
     );
 }
 
@@ -245,13 +254,16 @@ pub(crate) fn non_pep695_type_alias(checker: &Checker, stmt: &StmtAnnAssign) {
         .unique_by(|tvar| tvar.name)
         .collect::<Vec<_>>();
 
+    let range_with_parentheses =
+        parenthesized_range(value.into(), stmt.into(), checker.tokens()).unwrap_or(value.range());
+
     create_diagnostic(
         checker,
         stmt.into(),
         name,
-        value,
         &vars,
         TypeAliasKind::TypeAlias,
+        &checker.source()[range_with_parentheses],
     );
 }
 
@@ -260,9 +272,9 @@ fn create_diagnostic(
     checker: &Checker,
     stmt: StmtRef,
     name: &Name,
-    value: &Expr,
     type_vars: &[TypeVar],
     type_alias_kind: TypeAliasKind,
+    value_source: &str,
 ) {
     if type_vars.iter().any(TypeVar::has_unsupported_restriction) {
         return;
@@ -282,15 +294,10 @@ fn create_diagnostic(
     }
 
     let source = checker.source();
-    let tokens = checker.tokens();
-
-    let range_with_parentheses =
-        parenthesized_range(value.into(), stmt.into(), tokens).unwrap_or(value.range());
 
     let content = format!(
-        "type {name}{type_params} = {value}",
+        "type {name}{type_params} = {value_source}",
         type_params = DisplayTypeVars { type_vars, source },
-        value = &source[range_with_parentheses]
     );
     let edit = Edit::range_replacement(content, stmt.range());
 

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py.snap

@@ -443,3 +443,59 @@ help: Use the `type` keyword
 123 |     # comment1
     |
 note: This is an unsafe fix and may change runtime behavior
+
+UP040 [*] Type alias `LongAlias` uses `TypeAliasType` assignment instead of the `type` keyword
+   --> UP040.py:137:1
+    |
+136 |   # Preserve parentheses around a multiline union in a TypeAliasType call.
+137 | / LongAlias = TypeAliasType(
+138 | |     "LongAlias",
+139 | |     int
+140 | |     | str
+141 | |     | float
+142 | | )
+    | |_^
+143 |
+144 |   # An already parenthesized value keeps its own parentheses.
+    |
+help: Use the `type` keyword
+    |
+136 | # Preserve parentheses around a multiline union in a TypeAliasType call.
+    - LongAlias = TypeAliasType(
+    -     "LongAlias",
+    -     int
+137 + type LongAlias = (int
+138 |     | str
+    -     | float
+    - )
+139 +     | float)
+140 |
+    |
+note: This is an unsafe fix and may change runtime behavior
+
+UP040 [*] Type alias `ParenAlias` uses `TypeAliasType` assignment instead of the `type` keyword
+   --> UP040.py:145:1
+    |
+144 |   # An already parenthesized value keeps its own parentheses.
+145 | / ParenAlias = TypeAliasType(
+146 | |     "ParenAlias",
+147 | |     (
+148 | |         int
+149 | |         | str
+150 | |     ),
+151 | | )
+    | |_^
+help: Use the `type` keyword
+    |
+144 | # An already parenthesized value keeps its own parentheses.
+    - ParenAlias = TypeAliasType(
+    -     "ParenAlias",
+    -     (
+145 + type ParenAlias = (
+146 |         int
+147 |         | str
+    -     ),
+    - )
+148 +     )
+    |
+note: This is an unsafe fix and may change runtime behavior

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py__preview_diff.snap

@@ -57,10 +57,14 @@ UP040 [*] Type alias `DefaultList` uses `TypeAlias` annotation instead of the `t
 133 | T_default = TypeVar("T_default", default=int)
 134 | DefaultList: TypeAlias = list[T_default]
     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+135 |
+136 | # Preserve parentheses around a multiline union in a TypeAliasType call.
+    |
 help: Use the `type` keyword
     |
 133 | T_default = TypeVar("T_default", default=int)
     - DefaultList: TypeAlias = list[T_default]
 134 + type DefaultList[T_default = int] = list[T_default]
+135 |
     |
 note: This is an unsafe fix and may change runtime behavior