Add support for simple generic type variables to UP040

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

@@ -6,11 +6,26 @@
 x: TypeAlias = int
 
 
-# UP040 with generics (todo)
+# UP040 simple generic
 T = typing.TypeVar["T"]
 x: typing.TypeAlias = list[T]
 
 
+# UP040 bounded generic (todo)
+T = typing.TypeVar("T", bound=int)
+x: typing.TypeAlias = list[T]
+
+T = typing.TypeVar("T", int, str)
+x: typing.TypeAlias = list[T]
+
+# UP040 contravariant generic (todo)
+T = typing.TypeVar("T", contravariant=True)
+x: typing.TypeAlias = list[T]
+
+# UP040 covariant generic (todo)
+T = typing.TypeVar("T", covariant=True)
+x: typing.TypeAlias = list[T]
+
 # OK
 x: TypeAlias
 x: int = 1

crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs

@@ -1,4 +1,9 @@
-use ruff_python_ast::{Expr, ExprName, Ranged, Stmt, StmtAnnAssign, StmtTypeAlias};
+use ruff_python_ast::{
+    visitor::{self, Visitor},
+    Expr, ExprConstant, ExprName, ExprSubscript, Identifier, Ranged, Stmt, StmtAnnAssign,
+    StmtAssign, StmtTypeAlias, TypeParam, TypeParamTypeVar,
+};
+use ruff_python_semantic::SemanticModel;
 
 use crate::{registry::AsRule, settings::types::PythonVersion};
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
@@ -75,15 +80,78 @@ pub(crate) fn non_pep695_type_alias(checker: &mut Checker, stmt: &StmtAnnAssign)
     //              as type params instead
     let mut diagnostic = Diagnostic::new(NonPEP695TypeAlias { name: name.clone() }, stmt.range());
     if checker.patch(diagnostic.kind.rule()) {
+        let mut visitor = TypeVarNameVisitor {
+            names: vec![],
+            semantic: checker.semantic(),
+        };
+        visitor.visit_expr(value);
+
+        let type_params = if visitor.names.is_empty() {
+            None
+        } else {
+            Some(ruff_python_ast::TypeParams {
+                range: TextRange::default(),
+                type_params: visitor
+                    .names
+                    .iter()
+                    .map(|name| {
+                        TypeParam::TypeVar(TypeParamTypeVar {
+                            range: TextRange::default(),
+                            name: Identifier::new(name.id.clone(), TextRange::default()),
+                            bound: None,
+                        })
+                    })
+                    .collect(),
+            })
+        };
+
         diagnostic.set_fix(Fix::automatic(Edit::range_replacement(
             checker.generator().stmt(&Stmt::from(StmtTypeAlias {
                 range: TextRange::default(),
                 name: target.clone(),
-                type_params: None,
+                type_params,
                 value: value.clone(),
             })),
             stmt.range(),
         )));
     }
     checker.diagnostics.push(diagnostic);
 }
+
+struct TypeVarNameVisitor<'a> {
+    names: Vec<&'a ExprName>,
+    semantic: &'a SemanticModel<'a>,
+}
+
+impl<'a> Visitor<'a> for TypeVarNameVisitor<'a> {
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        match expr {
+            Expr::Name(name) if name.ctx.is_load() => {
+                let Some(Some(Stmt::Assign(StmtAssign { value, .. }))) =
+                    self.semantic
+                        .scope()
+                        .get(name.id.as_str())
+                        .map(|binding_id| {
+                            self.semantic
+                                .binding(binding_id)
+                                .source
+                                .map(|node_id| self.semantic.stmts[node_id])
+                        }) else {
+                            return;
+                        };
+
+                // Only support type variables declared as TypeVar['<name>'] for now
+                // Type variables declared with `TypeVar(<name>, ...)` can include more complex features
+                // like bounds and variance
+                let Expr::Subscript(ExprSubscript {value: ref subscript_value, .. })= value.as_ref() else {
+                    return;
+                };
+
+                if self.semantic.match_typing_expr(subscript_value, "TypeVar") {
+                    self.names.push(name);
+                }
+            }
+            _ => visitor::walk_expr(self, expr),
+        }
+    }
+}

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap

@@ -37,11 +37,11 @@ UP040.py:6:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of th
   6 |+type x = int
 7 7 | 
 8 8 | 
-9 9 | # UP040 with generics (todo)
+9 9 | # UP040 simple generic
 
 UP040.py:11:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword
    |
- 9 | # UP040 with generics (todo)
+ 9 | # UP040 simple generic
 10 | T = typing.TypeVar["T"]
 11 | x: typing.TypeAlias = list[T]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040
@@ -50,12 +50,95 @@ UP040.py:11:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t
 
 ℹ Fix
 8  8  | 
-9  9  | # UP040 with generics (todo)
+9  9  | # UP040 simple generic
 10 10 | T = typing.TypeVar["T"]
 11    |-x: typing.TypeAlias = list[T]
-   11 |+type x = list[T]
+   11 |+type x[T] = list[T]
 12 12 | 
 13 13 | 
-14 14 | # OK
+14 14 | # UP040 bounded generic (todo)
+
+UP040.py:16:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword
+   |
+14 | # UP040 bounded generic (todo)
+15 | T = typing.TypeVar("T", bound=int)
+16 | x: typing.TypeAlias = list[T]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040
+17 | 
+18 | T = typing.TypeVar("T", int, str)
+   |
+   = help: Use the `type` keyword
+
+ℹ Fix
+13 13 | 
+14 14 | # UP040 bounded generic (todo)
+15 15 | T = typing.TypeVar("T", bound=int)
+16    |-x: typing.TypeAlias = list[T]
+   16 |+type x = list[T]
+17 17 | 
+18 18 | T = typing.TypeVar("T", int, str)
+19 19 | x: typing.TypeAlias = list[T]
+
+UP040.py:19:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword
+   |
+18 | T = typing.TypeVar("T", int, str)
+19 | x: typing.TypeAlias = list[T]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040
+20 | 
+21 | # UP040 contravariant generic (todo)
+   |
+   = help: Use the `type` keyword
+
+ℹ Fix
+16 16 | x: typing.TypeAlias = list[T]
+17 17 | 
+18 18 | T = typing.TypeVar("T", int, str)
+19    |-x: typing.TypeAlias = list[T]
+   19 |+type x = list[T]
+20 20 | 
+21 21 | # UP040 contravariant generic (todo)
+22 22 | T = typing.TypeVar("T", contravariant=True)
+
+UP040.py:23:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword
+   |
+21 | # UP040 contravariant generic (todo)
+22 | T = typing.TypeVar("T", contravariant=True)
+23 | x: typing.TypeAlias = list[T]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040
+24 | 
+25 | # UP040 covariant generic (todo)
+   |
+   = help: Use the `type` keyword
+
+ℹ Fix
+20 20 | 
+21 21 | # UP040 contravariant generic (todo)
+22 22 | T = typing.TypeVar("T", contravariant=True)
+23    |-x: typing.TypeAlias = list[T]
+   23 |+type x = list[T]
+24 24 | 
+25 25 | # UP040 covariant generic (todo)
+26 26 | T = typing.TypeVar("T", covariant=True)
+
+UP040.py:27:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword
+   |
+25 | # UP040 covariant generic (todo)
+26 | T = typing.TypeVar("T", covariant=True)
+27 | x: typing.TypeAlias = list[T]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040
+28 | 
+29 | # OK
+   |
+   = help: Use the `type` keyword
+
+ℹ Fix
+24 24 | 
+25 25 | # UP040 covariant generic (todo)
+26 26 | T = typing.TypeVar("T", covariant=True)
+27    |-x: typing.TypeAlias = list[T]
+   27 |+type x = list[T]
+28 28 | 
+29 29 | # OK
+30 30 | x: TypeAlias