diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py index 49fb6f7c84673..8595dd8848145 100644 --- a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py +++ b/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] diff --git a/crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs b/crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs index 3adb84f7f38ac..08a78b31595f0 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs @@ -1,4 +1,6 @@ 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}, @@ -6,13 +8,10 @@ use ruff_python_ast::{ 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>, +} + struct TypeVarReferenceVisitor<'a> { - names: Vec<&'a ExprName>, + vars: Vec>, 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 }); } } _ => {} diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP040.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP040.py.snap index 33d0e3b165f09..f5c3dc9625cfc 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP040.py.snap +++ b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP040.py.snap @@ -69,7 +69,7 @@ UP040.py:14:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 14 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 15 | -16 | # UP040 bounded generic (todo) +16 | # UP040 bounded generic | = help: Use the `type` keyword @@ -80,153 +80,154 @@ UP040.py:14:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 14 |-x: typing.TypeAlias = list[T] 14 |+type x[T] = list[T] 15 15 | -16 16 | # UP040 bounded generic (todo) +16 16 | # UP040 bounded generic 17 17 | T = typing.TypeVar("T", bound=int) UP040.py:18:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword | -16 | # UP040 bounded generic (todo) +16 | # UP040 bounded generic 17 | T = typing.TypeVar("T", bound=int) 18 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 19 | -20 | T = typing.TypeVar("T", int, str) +20 | # UP040 constrained generic | = help: Use the `type` keyword ℹ Fix 15 15 | -16 16 | # UP040 bounded generic (todo) +16 16 | # UP040 bounded generic 17 17 | T = typing.TypeVar("T", bound=int) 18 |-x: typing.TypeAlias = list[T] - 18 |+type x = list[T] + 18 |+type x[T: int] = list[T] 19 19 | -20 20 | T = typing.TypeVar("T", int, str) -21 21 | x: typing.TypeAlias = list[T] +20 20 | # UP040 constrained generic +21 21 | T = typing.TypeVar("T", int, str) -UP040.py:21:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword +UP040.py:22:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword | -20 | T = typing.TypeVar("T", int, str) -21 | x: typing.TypeAlias = list[T] +20 | # UP040 constrained generic +21 | T = typing.TypeVar("T", int, str) +22 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -22 | -23 | # UP040 contravariant generic (todo) +23 | +24 | # UP040 contravariant generic | = help: Use the `type` keyword ℹ Fix -18 18 | x: typing.TypeAlias = list[T] 19 19 | -20 20 | T = typing.TypeVar("T", int, str) -21 |-x: typing.TypeAlias = list[T] - 21 |+type x = list[T] -22 22 | -23 23 | # UP040 contravariant generic (todo) -24 24 | T = typing.TypeVar("T", contravariant=True) - -UP040.py:25:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword - | -23 | # UP040 contravariant generic (todo) -24 | T = typing.TypeVar("T", contravariant=True) -25 | x: typing.TypeAlias = list[T] +20 20 | # UP040 constrained generic +21 21 | T = typing.TypeVar("T", int, str) +22 |-x: typing.TypeAlias = list[T] + 22 |+type x[T: (int, str)] = list[T] +23 23 | +24 24 | # UP040 contravariant generic +25 25 | T = typing.TypeVar("T", contravariant=True) + +UP040.py:26:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword + | +24 | # UP040 contravariant generic +25 | T = typing.TypeVar("T", contravariant=True) +26 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -26 | -27 | # UP040 covariant generic (todo) +27 | +28 | # UP040 covariant generic | = help: Use the `type` keyword ℹ Fix -22 22 | -23 23 | # UP040 contravariant generic (todo) -24 24 | T = typing.TypeVar("T", contravariant=True) -25 |-x: typing.TypeAlias = list[T] - 25 |+type x = list[T] -26 26 | -27 27 | # UP040 covariant generic (todo) -28 28 | T = typing.TypeVar("T", covariant=True) - -UP040.py:29:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword - | -27 | # UP040 covariant generic (todo) -28 | T = typing.TypeVar("T", covariant=True) -29 | x: typing.TypeAlias = list[T] +23 23 | +24 24 | # UP040 contravariant generic +25 25 | T = typing.TypeVar("T", contravariant=True) +26 |-x: typing.TypeAlias = list[T] + 26 |+type x[T] = list[T] +27 27 | +28 28 | # UP040 covariant generic +29 29 | T = typing.TypeVar("T", covariant=True) + +UP040.py:30:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword + | +28 | # UP040 covariant generic +29 | T = typing.TypeVar("T", covariant=True) +30 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -30 | -31 | # UP040 in class scope +31 | +32 | # UP040 in class scope | = help: Use the `type` keyword ℹ Fix -26 26 | -27 27 | # UP040 covariant generic (todo) -28 28 | T = typing.TypeVar("T", covariant=True) -29 |-x: typing.TypeAlias = list[T] - 29 |+type x = list[T] -30 30 | -31 31 | # UP040 in class scope -32 32 | T = typing.TypeVar["T"] - -UP040.py:35:5: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword - | -33 | class Foo: -34 | # reference to global variable -35 | x: typing.TypeAlias = list[T] +27 27 | +28 28 | # UP040 covariant generic +29 29 | T = typing.TypeVar("T", covariant=True) +30 |-x: typing.TypeAlias = list[T] + 30 |+type x[T] = list[T] +31 31 | +32 32 | # UP040 in class scope +33 33 | T = typing.TypeVar["T"] + +UP040.py:36:5: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword + | +34 | class Foo: +35 | # reference to global variable +36 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -36 | -37 | # reference to class variable +37 | +38 | # reference to class variable | = help: Use the `type` keyword ℹ Fix -32 32 | T = typing.TypeVar["T"] -33 33 | class Foo: -34 34 | # reference to global variable -35 |- x: typing.TypeAlias = list[T] - 35 |+ type x[T] = list[T] -36 36 | -37 37 | # reference to class variable -38 38 | TCLS = typing.TypeVar["TCLS"] - -UP040.py:39:5: UP040 [*] Type alias `y` uses `TypeAlias` annotation instead of the `type` keyword - | -37 | # reference to class variable -38 | TCLS = typing.TypeVar["TCLS"] -39 | y: typing.TypeAlias = list[TCLS] +33 33 | T = typing.TypeVar["T"] +34 34 | class Foo: +35 35 | # reference to global variable +36 |- x: typing.TypeAlias = list[T] + 36 |+ type x[T] = list[T] +37 37 | +38 38 | # reference to class variable +39 39 | TCLS = typing.TypeVar["TCLS"] + +UP040.py:40:5: UP040 [*] Type alias `y` uses `TypeAlias` annotation instead of the `type` keyword + | +38 | # reference to class variable +39 | TCLS = typing.TypeVar["TCLS"] +40 | y: typing.TypeAlias = list[TCLS] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -40 | -41 | # UP040 wont add generics in fix +41 | +42 | # UP040 wont add generics in fix | = help: Use the `type` keyword ℹ Fix -36 36 | -37 37 | # reference to class variable -38 38 | TCLS = typing.TypeVar["TCLS"] -39 |- y: typing.TypeAlias = list[TCLS] - 39 |+ type y[TCLS] = list[TCLS] -40 40 | -41 41 | # UP040 wont add generics in fix -42 42 | T = typing.TypeVar(*args) - -UP040.py:43:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword - | -41 | # UP040 wont add generics in fix -42 | T = typing.TypeVar(*args) -43 | x: typing.TypeAlias = list[T] +37 37 | +38 38 | # reference to class variable +39 39 | TCLS = typing.TypeVar["TCLS"] +40 |- y: typing.TypeAlias = list[TCLS] + 40 |+ type y[TCLS] = list[TCLS] +41 41 | +42 42 | # UP040 wont add generics in fix +43 43 | T = typing.TypeVar(*args) + +UP040.py:44:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword + | +42 | # UP040 wont add generics in fix +43 | T = typing.TypeVar(*args) +44 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -44 | -45 | # OK +45 | +46 | # OK | = help: Use the `type` keyword ℹ Fix -40 40 | -41 41 | # UP040 wont add generics in fix -42 42 | T = typing.TypeVar(*args) -43 |-x: typing.TypeAlias = list[T] - 43 |+type x = list[T] -44 44 | -45 45 | # OK -46 46 | x: TypeAlias +41 41 | +42 42 | # UP040 wont add generics in fix +43 43 | T = typing.TypeVar(*args) +44 |-x: typing.TypeAlias = list[T] + 44 |+type x = list[T] +45 45 | +46 46 | # OK +47 47 | x: TypeAlias