diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py index 69854a8af75f4..49fb6f7c84673 100644 --- a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py +++ b/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py @@ -5,11 +5,42 @@ x: typing.TypeAlias = int x: TypeAlias = int - -# UP040 with generics (todo) +# UP040 simple generic T = typing.TypeVar["T"] x: typing.TypeAlias = list[T] +# UP040 call style 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] + +# UP040 in class scope +T = typing.TypeVar["T"] +class Foo: + # reference to global variable + x: typing.TypeAlias = list[T] + + # reference to class variable + TCLS = typing.TypeVar["TCLS"] + y: typing.TypeAlias = list[TCLS] + +# UP040 wont add generics in fix +T = typing.TypeVar(*args) +x: typing.TypeAlias = list[T] # OK x: TypeAlias 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 64d00f6a80cf9..3db1ec67df8b0 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,11 @@ -use ruff_python_ast::{Expr, ExprName, Ranged, Stmt, StmtAnnAssign, StmtTypeAlias}; +use ast::{Constant, ExprCall, ExprConstant}; +use ruff_python_ast::{ + self as ast, + visitor::{self, Visitor}, + Expr, 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,11 +82,36 @@ 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 = TypeVarReferenceVisitor { + names: vec![], + semantic: checker.semantic(), + }; + visitor.visit_expr(value); + + let type_params = if visitor.names.is_empty() { + None + } else { + Some(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(), @@ -87,3 +119,62 @@ pub(crate) fn non_pep695_type_alias(checker: &mut Checker, stmt: &StmtAnnAssign) } checker.diagnostics.push(diagnostic); } + +struct TypeVarReferenceVisitor<'a> { + names: Vec<&'a ExprName>, + semantic: &'a SemanticModel<'a>, +} + +/// Recursively collects the names of type variable references present in an expression. +impl<'a> Visitor<'a> for TypeVarReferenceVisitor<'a> { + fn visit_expr(&mut self, expr: &'a Expr) { + match expr { + Expr::Name(name) if name.ctx.is_load() => { + let Some(Stmt::Assign(StmtAssign { value, .. })) = + self.semantic.lookup_symbol(name.id.as_str()) + .and_then(|binding_id| { + self.semantic + .binding(binding_id) + .source + .map(|node_id| self.semantic.statement(node_id)) + }) else { + return; + }; + + match value.as_ref() { + Expr::Subscript(ExprSubscript { + value: ref subscript_value, + .. + }) => { + if self.semantic.match_typing_expr(subscript_value, "TypeVar") { + self.names.push(name); + } + } + 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, + Expr::Constant(ExprConstant { + value: Constant::Str(_), + .. + }) + ) + }) + && arguments.keywords.is_empty() + { + self.names.push(name); + } + } + _ => {} + } + } + _ => visitor::walk_expr(self, expr), + } + } +} diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap index c46a4b8879b8f..33d0e3b165f09 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap +++ b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap @@ -18,7 +18,7 @@ UP040.py:5:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of th 5 |+type x = int 6 6 | x: TypeAlias = int 7 7 | -8 8 | +8 8 | # UP040 simple generic UP040.py:6:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword | @@ -26,6 +26,8 @@ UP040.py:6:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of th 5 | x: typing.TypeAlias = int 6 | x: TypeAlias = int | ^^^^^^^^^^^^^^^^^^ UP040 +7 | +8 | # UP040 simple generic | = help: Use the `type` keyword @@ -36,26 +38,195 @@ UP040.py:6:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of th 6 |-x: TypeAlias = int 6 |+type x = int 7 7 | -8 8 | -9 9 | # UP040 with generics (todo) +8 8 | # UP040 simple generic +9 9 | T = typing.TypeVar["T"] -UP040.py:11:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword +UP040.py:10:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword | - 9 | # UP040 with generics (todo) -10 | T = typing.TypeVar["T"] -11 | x: typing.TypeAlias = list[T] + 8 | # UP040 simple generic + 9 | T = typing.TypeVar["T"] +10 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +11 | +12 | # UP040 call style generic | = help: Use the `type` keyword ℹ Fix -8 8 | -9 9 | # UP040 with generics (todo) -10 10 | T = typing.TypeVar["T"] -11 |-x: typing.TypeAlias = list[T] - 11 |+type x = list[T] -12 12 | -13 13 | -14 14 | # OK +7 7 | +8 8 | # UP040 simple generic +9 9 | T = typing.TypeVar["T"] +10 |-x: typing.TypeAlias = list[T] + 10 |+type x[T] = list[T] +11 11 | +12 12 | # UP040 call style generic +13 13 | T = typing.TypeVar("T") + +UP040.py:14:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword + | +12 | # UP040 call style generic +13 | T = typing.TypeVar("T") +14 | x: typing.TypeAlias = list[T] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +15 | +16 | # UP040 bounded generic (todo) + | + = help: Use the `type` keyword + +ℹ Fix +11 11 | +12 12 | # UP040 call style generic +13 13 | T = typing.TypeVar("T") +14 |-x: typing.TypeAlias = list[T] + 14 |+type x[T] = list[T] +15 15 | +16 16 | # UP040 bounded generic (todo) +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) +17 | T = typing.TypeVar("T", bound=int) +18 | x: typing.TypeAlias = list[T] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +19 | +20 | T = typing.TypeVar("T", int, str) + | + = help: Use the `type` keyword + +ℹ Fix +15 15 | +16 16 | # UP040 bounded generic (todo) +17 17 | T = typing.TypeVar("T", bound=int) +18 |-x: typing.TypeAlias = list[T] + 18 |+type x = list[T] +19 19 | +20 20 | T = typing.TypeVar("T", int, str) +21 21 | x: typing.TypeAlias = list[T] + +UP040.py:21: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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +22 | +23 | # UP040 contravariant generic (todo) + | + = 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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +26 | +27 | # UP040 covariant generic (todo) + | + = 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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +30 | +31 | # 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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +36 | +37 | # 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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +40 | +41 | # 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] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 +44 | +45 | # 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