Skip to content

Commit

Permalink
Add support for TypeVar("..") style generics
Browse files Browse the repository at this point in the history
  • Loading branch information
@zanieb
zanieb committed Aug 7, 2023
1 parent 75584d8 commit 2fffafb
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 47 deletions.
7 changes: 4 additions & 3 deletions crates/ruff/resources/test/fixtures/pyupgrade/UP040.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand All @@ -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:
Expand All @@ -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
Expand Down
41 changes: 32 additions & 9 deletions crates/ruff/src/rules/pyupgrade/rules/use_pep695_type_alias.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use ast::{Constant, ExprCall, ExprConstant};
use ruff_python_ast::{
self as ast,
visitor::{self, Visitor},
Expand Down Expand Up @@ -140,15 +141,37 @@ impl<'a> Visitor<'a> for TypeVarReferenceVisitor<'a> {
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);
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),
Expand Down
95 changes: 60 additions & 35 deletions ...rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__non_pep695_type_alias_py312.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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

0 comments on commit 2fffafb

Please sign in to comment.