From 2fffafbc441002924ba533698ff5fdf9f46c7dee Mon Sep 17 00:00:00 2001 From: Zanie Date: Mon, 7 Aug 2023 14:42:23 -0500 Subject: [PATCH] Add support for `TypeVar("..")` style generics --- .../test/fixtures/pyupgrade/UP040.py | 7 +- .../pyupgrade/rules/use_pep695_type_alias.rs | 41 ++++++-- ...e__tests__non_pep695_type_alias_py312.snap | 95 ++++++++++++------- 3 files changed, 96 insertions(+), 47 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py index c3c0b2eb32d5f7..49fb6f7c84673a 100644 --- a/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py +++ b/crates/ruff/resources/test/fixtures/pyupgrade/UP040.py @@ -9,7 +9,7 @@ T = typing.TypeVar["T"] x: typing.TypeAlias = list[T] -# UP040 call style generic (todo) +# UP040 call style generic T = typing.TypeVar("T") x: typing.TypeAlias = list[T] @@ -28,7 +28,6 @@ T = typing.TypeVar("T", covariant=True) x: typing.TypeAlias = list[T] - # UP040 in class scope T = typing.TypeVar["T"] class Foo: @@ -39,7 +38,9 @@ class Foo: 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 8087da1a6cd09e..82714b9d7fae22 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,3 +1,4 @@ +use ast::{Constant, ExprCall, ExprConstant}; use ruff_python_ast::{ self as ast, visitor::{self, Visitor}, @@ -140,15 +141,37 @@ impl<'a> Visitor<'a> for TypeVarReferenceVisitor<'a> { return; }; - // Only support type variables declared as TypeVar[''] for now - // Type variables declared with `TypeVar(, ...)` 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); + 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.len() == 0 + { + 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 d82eb8caded3d5..33d0e3b165f096 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 @@ -48,7 +48,7 @@ UP040.py:10:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 10 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 11 | -12 | # UP040 call style generic (todo) +12 | # UP040 call style generic | = help: Use the `type` keyword @@ -59,12 +59,12 @@ UP040.py:10:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 10 |-x: typing.TypeAlias = list[T] 10 |+type x[T] = list[T] 11 11 | -12 12 | # UP040 call style generic (todo) +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 (todo) +12 | # UP040 call style generic 13 | T = typing.TypeVar("T") 14 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 @@ -75,10 +75,10 @@ UP040.py:14:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t ℹ Fix 11 11 | -12 12 | # UP040 call style generic (todo) +12 12 | # UP040 call style generic 13 13 | T = typing.TypeVar("T") 14 |-x: typing.TypeAlias = list[T] - 14 |+type x = list[T] + 14 |+type x[T] = list[T] 15 15 | 16 16 | # UP040 bounded generic (todo) 17 17 | T = typing.TypeVar("T", bound=int) @@ -151,6 +151,8 @@ UP040.py:29:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 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 @@ -161,47 +163,70 @@ UP040.py:29:1: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of t 29 |-x: typing.TypeAlias = list[T] 29 |+type x = list[T] 30 30 | -31 31 | -32 32 | # UP040 in class scope +31 31 | # UP040 in class scope +32 32 | T = typing.TypeVar["T"] -UP040.py:36:5: UP040 [*] Type alias `x` uses `TypeAlias` annotation instead of the `type` keyword +UP040.py:35: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] +33 | class Foo: +34 | # reference to global variable +35 | x: typing.TypeAlias = list[T] | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP040 -37 | -38 | # reference to class variable +36 | +37 | # reference to class variable | = help: Use the `type` keyword ℹ Fix -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] +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 -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 | -43 43 | +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