Add support for bounds, constraints, and explicit variance on generic type variables to UP040

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

@@ -13,18 +13,19 @@
 T = typing.TypeVar("T")
 x: typing.TypeAlias = list[T]
 
-# UP040 bounded generic (todo)
+# UP040 bounded generic
 T = typing.TypeVar("T", bound=int)
 x: typing.TypeAlias = list[T]
 
+# UP040 constrained generic
 T = typing.TypeVar("T", int, str)
 x: typing.TypeAlias = list[T]
 
-# UP040 contravariant generic (todo)
+# UP040 contravariant generic
 T = typing.TypeVar("T", contravariant=True)
 x: typing.TypeAlias = list[T]
 
-# UP040 covariant generic (todo)
+# UP040 covariant generic
 T = typing.TypeVar("T", covariant=True)
 x: typing.TypeAlias = list[T]
 

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

@@ -1,18 +1,17 @@
 use ast::{Constant, ExprCall, ExprConstant};
+use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
+use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{
     self as ast,
     visitor::{self, Visitor},
     Expr, ExprName, ExprSubscript, Identifier, 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};
-use ruff_macros::{derive_message_formats, violation};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
+use crate::{registry::AsRule, settings::types::PythonVersion};
 
 /// ## What it does
 /// Checks for use of `TypeAlias` annotation for declaring type aliases.
@@ -83,24 +82,36 @@ pub(crate) fn non_pep695_type_alias(checker: &mut Checker, stmt: &StmtAnnAssign)
     let mut diagnostic = Diagnostic::new(NonPEP695TypeAlias { name: name.clone() }, stmt.range());
     if checker.patch(diagnostic.kind.rule()) {
         let mut visitor = TypeVarReferenceVisitor {
-            names: vec![],
+            vars: vec![],
             semantic: checker.semantic(),
         };
         visitor.visit_expr(value);
 
-        let type_params = if visitor.names.is_empty() {
+        let type_params = if visitor.vars.is_empty() {
             None
         } else {
             Some(ast::TypeParams {
                 range: TextRange::default(),
                 type_params: visitor
-                    .names
-                    .iter()
-                    .map(|name| {
+                    .vars
+                    .into_iter()
+                    .map(|TypeVar { name, restriction }| {
                         TypeParam::TypeVar(TypeParamTypeVar {
                             range: TextRange::default(),
                             name: Identifier::new(name.id.clone(), TextRange::default()),
-                            bound: None,
+                            bound: match restriction {
+                                Some(TypeVarRestriction::Bound(bound)) => {
+                                    Some(Box::new(bound.clone()))
+                                }
+                                Some(TypeVarRestriction::Constraint(constraints)) => {
+                                    Some(Box::new(Expr::Tuple(ast::ExprTuple {
+                                        range: TextRange::default(),
+                                        elts: constraints.into_iter().cloned().collect(),
+                                        ctx: ast::ExprContext::Load,
+                                    })))
+                                }
+                                None => None,
+                            },
                         })
                     })
                     .collect(),
@@ -120,8 +131,22 @@ pub(crate) fn non_pep695_type_alias(checker: &mut Checker, stmt: &StmtAnnAssign)
     checker.diagnostics.push(diagnostic);
 }
 
+#[derive(Debug)]
+enum TypeVarRestriction<'a> {
+    /// A type variable with a bound, e.g., `TypeVar("T", bound=int)`.
+    Bound(&'a Expr),
+    /// A type variable with constraints, e.g., `TypeVar("T", int, str)`.
+    Constraint(Vec<&'a Expr>),
+}
+
+#[derive(Debug)]
+struct TypeVar<'a> {
+    name: &'a ExprName,
+    restriction: Option<TypeVarRestriction<'a>>,
+}
+
 struct TypeVarReferenceVisitor<'a> {
-    names: Vec<&'a ExprName>,
+    vars: Vec<TypeVar<'a>>,
     semantic: &'a SemanticModel<'a>,
 }
 
@@ -149,16 +174,16 @@ impl<'a> Visitor<'a> for TypeVarReferenceVisitor<'a> {
                         ..
                     }) => {
                         if self.semantic.match_typing_expr(subscript_value, "TypeVar") {
-                            self.names.push(name);
+                            self.vars.push(TypeVar {
+                                name,
+                                restriction: None,
+                            });
                         }
                     }
                     Expr::Call(ExprCall {
                         func, arguments, ..
                     }) => {
-                        // TODO(zanieb): Add support for bounds and variance declarations
-                        //               for now this only supports `TypeVar("...")`
                         if self.semantic.match_typing_expr(func, "TypeVar")
-                            && arguments.args.len() == 1
                             && arguments.args.first().is_some_and(|arg| {
                                 matches!(
                                     arg,
@@ -168,9 +193,18 @@ impl<'a> Visitor<'a> for TypeVarReferenceVisitor<'a> {
                                     })
                                 )
                             })
-                            && arguments.keywords.is_empty()
                         {
-                            self.names.push(name);
+                            let restriction = if let Some(bound) = arguments.find_keyword("bound") {
+                                Some(TypeVarRestriction::Bound(&bound.value))
+                            } else if arguments.args.len() > 1 {
+                                Some(TypeVarRestriction::Constraint(
+                                    arguments.args.iter().skip(1).collect(),
+                                ))
+                            } else {
+                                None
+                            };
+
+                            self.vars.push(TypeVar { name, restriction });
                         }
                     }
                     _ => {}