From 06e45034498831edab509136eb39f2a79bcb368b Mon Sep 17 00:00:00 2001
From: Alex Waygood <alex.waygood@gmail.com>
Date: Fri, 7 Feb 2025 16:59:18 +0000
Subject: [PATCH] more tests + use `StmtRef`

---
 .../test/fixtures/pyupgrade/UP040.py          | 10 +++++++
 .../rules/pep695/non_pep695_type_alias.rs     | 10 ++++---
 ...er__rules__pyupgrade__tests__UP040.py.snap | 29 +++++++++++++++++++
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP040.py b/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP040.py
index 11fb5faea9717..f6797167e1f18 100644
--- a/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP040.py
+++ b/crates/ruff_linter/resources/test/fixtures/pyupgrade/UP040.py
@@ -118,3 +118,13 @@ class Foo:
     int
     | str
 )
+
+T: TypeAlias = ( # comment0
+    # comment1
+    int  # comment2
+    # comment3
+    | # comment4
+    # comment5
+    str  # comment6
+    # comment7
+) # comment8
diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/pep695/non_pep695_type_alias.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/pep695/non_pep695_type_alias.rs
index 1622268ad384f..ff9215f962535 100644
--- a/crates/ruff_linter/src/rules/pyupgrade/rules/pep695/non_pep695_type_alias.rs
+++ b/crates/ruff_linter/src/rules/pyupgrade/rules/pep695/non_pep695_type_alias.rs
@@ -5,7 +5,7 @@ use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::name::Name;
 use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::visitor::Visitor;
-use ruff_python_ast::{AnyNodeRef, Expr, ExprCall, ExprName, Keyword, StmtAnnAssign, StmtAssign};
+use ruff_python_ast::{Expr, ExprCall, ExprName, Keyword, StmtAnnAssign, StmtAssign, StmtRef};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -262,7 +262,7 @@ pub(crate) fn non_pep695_type_alias(checker: &Checker, stmt: &StmtAnnAssign) {
 /// Generate a [`Diagnostic`] for a non-PEP 695 type alias or type alias type.
 fn create_diagnostic(
     checker: &Checker,
-    stmt: AnyNodeRef,
+    stmt: StmtRef,
     name: &Name,
     value: &Expr,
     type_vars: &[TypeVar],
@@ -270,11 +270,13 @@ fn create_diagnostic(
     type_alias_kind: TypeAliasKind,
 ) -> Diagnostic {
     let source = checker.source();
+    let range_with_parentheses =
+        parenthesized_range(value.into(), stmt.into(), checker.comment_ranges(), source)
+            .unwrap_or(value.range());
     let content = format!(
         "type {name}{type_params} = {value}",
         type_params = DisplayTypeVars { type_vars, source },
-        value = &source[parenthesized_range(value.into(), stmt, checker.comment_ranges(), source)
-            .unwrap_or(value.range())]
+        value = &source[range_with_parentheses]
     );
     let edit = Edit::range_replacement(content, stmt.range());
     Diagnostic::new(
diff --git a/crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py.snap b/crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py.snap
index 9a3fecd428f40..1c5818c218f29 100644
--- a/crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py.snap
+++ b/crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP040.py.snap
@@ -446,6 +446,8 @@ UP040.py:117:1: UP040 [*] Type alias `T` uses `TypeAlias` annotation instead of
 119 | |     | str
 120 | | )
     | |_^ UP040
+121 |
+122 |   T: TypeAlias = ( # comment0
     |
     = help: Use the `type` keyword
 
@@ -458,3 +460,30 @@ UP040.py:117:1: UP040 [*] Type alias `T` uses `TypeAlias` annotation instead of
 118 118 |     int
 119 119 |     | str
 120 120 | )
+
+UP040.py:122:1: UP040 [*] Type alias `T` uses `TypeAlias` annotation instead of the `type` keyword
+    |
+120 |   )
+121 |
+122 | / T: TypeAlias = ( # comment0
+123 | |     # comment1
+124 | |     int  # comment2
+125 | |     # comment3
+126 | |     | # comment4
+127 | |     # comment5
+128 | |     str  # comment6
+129 | |     # comment7
+130 | | ) # comment8
+    | |_^ UP040
+    |
+    = help: Use the `type` keyword
+
+ℹ Unsafe fix
+119 119 |     | str
+120 120 | )
+121 121 | 
+122     |-T: TypeAlias = ( # comment0
+    122 |+type T = ( # comment0
+123 123 |     # comment1
+124 124 |     int  # comment2
+125 125 |     # comment3