Skip to content

Commit

Permalink
Add support for simple generic type variables to UP040 (astral-sh#6314)
Browse files Browse the repository at this point in the history 
Extends astral-sh#6289 to support moving type variable usage in type aliases to
use PEP-695.

Does not remove the possibly unused type variable declaration.
Presumably this is handled by other rules, but is not working for me.

Does not handle type variables with bounds or variance declarations yet.

Part of astral-sh#4617
  • Loading branch information
@zanieb @durumu
zanieb authored and durumu committed Aug 12, 2023
1 parent b7f4743 commit 85f3f09
Show file tree
Hide file tree
Showing 3 changed files with 312 additions and 19 deletions.
35 changes: 33 additions & 2 deletions crates/ruff/resources/test/fixtures/pyupgrade/UP040.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
95 changes: 93 additions & 2 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,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};
Expand Down Expand Up @@ -75,15 +82,99 @@ 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(),
)));
}
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),
}
}
}
201 changes: 186 additions & 15 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 @@ -18,14 +18,16 @@ 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
|
4 | # UP040
5 | x: typing.TypeAlias = int
6 | x: TypeAlias = int
| ^^^^^^^^^^^^^^^^^^ UP040
7 |
8 | # UP040 simple generic
|
= help: Use the `type` keyword
Expand All @@ -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

0 comments on commit 85f3f09

Please sign in to comment.